Keras not accepting character `/` from build_factored_surrogate_posterior

[WillianFuks]

Given this brief code:

import pandas as pd
import tensorflow_probability as tfp
import tensorflow as tf

data = pd.read_csv('https://raw.githubusercontent.com/WillianFuks/tfcausalimpact/master/tests/fixtures/arma_data.csv')[['y']].astype('float32')
data.index = pd.date_range(start='2024-01-01', periods=len(data), freq='D')
obs = data.iloc[:70]
model = tfp.sts.Sum([tfp.sts.LocalLevel(observed_time_series=obs)])

optimizer = tf.optimizers.Adam(learning_rate=0.1)
variational_steps = 200
variational_posteriors = tfp.sts.build_factored_surrogate_posterior(model=model)

@tf.function()
def _run_vi():
    tfp.vi.fit_surrogate_posterior(
        target_log_prob_fn=model.joint_log_prob(
            observed_time_series=obs
        ),
        surrogate_posterior=variational_posteriors,
        optimizer=optimizer,
        num_steps=variational_steps
    )
    samples = variational_posteriors.sample(100)
    return samples, None
samples, _ = _run_vi()

Its execution raises:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[<ipython-input-24-6998a2ed04ce>](https://localhost:8080/#) in <cell line: 22>()
     20     samples = variational_posteriors.sample(100)
     21     return samples, None
---> 22 samples, _ = _run_vi()
     23 

12 frames
[/usr/local/lib/python3.10/dist-packages/keras/src/backend/common/variables.py](https://localhost:8080/#) in __init__(self, initializer, shape, dtype, trainable, name)
     17         name = name or auto_name(self.__class__.__name__)
     18         if not isinstance(name, str) or "/" in name:
---> 19             raise ValueError(
     20                 "Argument `name` must be a string and "
     21                 "cannot contain character `/`. "

ValueError: in user code:

    File "<ipython-input-21-6998a2ed04ce>", line 10, in _run_vi  *
        tfp.vi.fit_surrogate_posterior(
    File "/usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/vi/optimization.py", line 724, in fit_surrogate_posterior  **
        return minimize(
    File "/usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/math/minimize.py", line 616, in minimize
        _, traced_values = _minimize_common(
    File "/usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/math/minimize.py", line 156, in _minimize_common
        initial_optimizer_state) = optimizer_step_fn(
    File "/usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/math/minimize.py", line 435, in optimizer_step
        train_op = optimizer.apply_gradients(zip(grads, watched_variables))
    File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/base_optimizer.py", line 269, in apply_gradients
        self.apply(grads, trainable_variables)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/base_optimizer.py", line 308, in apply
        self.build(trainable_variables)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/adam.py", line 93, in build
        self.add_variable_from_reference(
    File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/tensorflow/optimizer.py", line 33, in add_variable_from_reference
        return super().add_variable_from_reference(
    File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/base_optimizer.py", line 205, in add_variable_from_reference
        return self.add_variable(
    File "/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/base_optimizer.py", line 184, in add_variable
        variable = backend.Variable(
    File "/usr/local/lib/python3.10/dist-packages/keras/src/backend/common/variables.py", line 19, in __init__
        raise ValueError(

    ValueError: Argument `name` must be a string and cannot contain character `/`. Received: name=build_factored_surrogate_posterior/loc_0_momentum

Apparently the integration with Keras raised conflicts on the naming patterns used in tfp.

Versions used: tensorflow==2.16.1 tensorflow-probability[tf]==0.24.0 tf-keras==2.16.0

[Llamrei]

Just FYI I believe this is an issue with keras - see here

[jburnim]

Does the error go away if you add:

import tf_keras

and replace tf.optimizers with tf_keras.optimizers?

In TF 2.16+, tf.optimizers refers to the Keras 3 version of the optimizer, which introduced this breaking change. But tf_keras.optimizers should give you the Keras 2 version of the optimizer (which is what tf.optimizers did in TF 2.15 and earlier).

[shilinng]

I had the same problem and replacing tf.optimizers with tf_keras.optimizers solved the problem. Will there be a fix here?

Does the error go away if you add:

import tf_keras

and replace tf.optimizers with tf_keras.optimizers?

In TF 2.16+, tf.optimizers refers to the Keras 3 version of the optimizer, which introduced this breaking change. But tf_keras.optimizers should give you the Keras 2 version of the optimizer (which is what tf.optimizers did in TF 2.15 and earlier).

[jburnim]

From the TFP 0.24.0 release notes -- https://github.com/tensorflow/probability/releases/tag/v0.24.0 :

NOTE: In TensorFlow 2.16+, tf.keras (and tf.initializers, tf.losses, and tf.optimizers) refers to Keras 3. TensorFlow Probability is not compatible with Keras 3 -- instead TFP is continuing to use Keras 2, which is now packaged as tf-keras and tf-keras-nightly and is imported as tf_keras.

We do not plan to update TensorFlow Probability to be compatible with Keras 3.

When TensorFlow moved from Keras 2 to Keras 3, it looks like it introduced some breaking changes in tf.optimizers, as well as tf.keras (and potentially tf.initializers and tf.losses). To avoid these changes and stay on Keras 2, users must switch from usingtf.optimizers to tf_keras.optimizers (and similarly for tf.initializers and tf.losses).