My Goal: Use the add_loss method inside a custom RNN cell (in graph execution mode) to add an input-dependent loss.

General Setup:

• Using Python 3.9
• Using TensorFlow 2.8 or 2.10
• Assuming import tensorflow as tf, I have a subclassed tf.keras.Model that uses a standard tf.keras.layers.RNN layer and a custom RNN cell (subclasses tf.keras.layers.Layer). Inside my custom RNN cell I call self.add_loss(*) in order to add an input-dependent loss.

Expected Result: When I call Model.fit(), the add_loss method is called for every batch and every timestep. The gradient computation step uses the added losses without raising an error.

Actual Result: When I call Model.fit(), an InaccessibleTensorError is raised during the gradient computation step, specifically when self.losses is called inside Model.train_step().

Exception has occurred: InaccessibleTensorError
<tf.Tensor 'foo_model/rnn/while/bar_cell/Sum_1:0' shape=() dtype=float32> is out of scope and cannot be used here. Use return values, explicit Python locals or TensorFlow collections to access it.
Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values for more information.

What I've tried:

• The error is not raised when using eager execution.
• The error is not raised when using graph execution and initializing the RNN layer with unrolled=True. Unfortunately this doesn't help me since my sequences can be long.
• Switching to the latest stable version of TensorFlow (2.10.0) does not fix the issue.
• After searching the web, Stack Overflow and issues/code on TensorFlow's GitHub, I'm completely stumped.

Minimum Reproducible Example

import pytest
import tensorflow as tf


class FooModel(tf.keras.Model):
    """A basic model for testing.

    Attributes:
        cell: The RNN cell layer.

    """

    def __init__(self, cell=None, **kwargs):
        """Initialize.

        Args:
            cell: A Keras layer.
            kwargs:  Additional key-word arguments.

        Raises:
            ValueError: If arguments are invalid.

        """
        super().__init__(**kwargs)

        # Assign layers.
        self.rnn = tf.keras.layers.RNN(cell, return_sequences=True)

    def call(self, inputs, training=None):
        """Call.

        Args:
            inputs: A dictionary of inputs.
            training (optional): Boolean indicating if training mode.

        """
        output = self.rnn(inputs, training=training)
        return output


class BarCell(tf.keras.layers.Layer):
    """RNN cell for testing."""
    def __init__(self, **kwargs):
        """Initialize.

        Args:

        """
        super(BarCell, self).__init__(**kwargs)

        # Satisfy RNNCell contract.
        self.state_size = [tf.TensorShape([1]),]

    def call(self, inputs, states, training=None):
        """Call."""
        output = tf.reduce_sum(inputs, axis=1) + tf.constant(1.0)
        self.add_loss(tf.reduce_sum(inputs))

        states_tplus1 = [states[0] + 1]
        return output, states_tplus1


@pytest.mark.parametrize(
    "is_eager", [True, False]
)
def test_rnn_fit_with_add_loss(is_eager):
    """Test fit method (triggering backprop)."""
    tf.config.run_functions_eagerly(is_eager)

    # Some dummy input formatted as a TF Dataset.
    n_example = 5
    x = tf.constant([
        [[1, 2, 3], [2, 0, 0], [3, 0, 0], [4, 3, 4]],
        [[1, 13, 8], [2, 0, 0], [3, 0, 0], [4, 13, 8]],
        [[1, 5, 6], [2, 8, 0], [3, 16, 0], [4, 5, 6]],
        [[1, 5, 12], [2, 14, 15], [3, 17, 18], [4, 5, 6]],
        [[1, 5, 6], [2, 14, 15], [3, 17, 18], [4, 5, 6]],
    ], dtype=tf.float32)
    y = tf.constant(
        [
            [[1], [2], [1], [2]],
            [[10], [2], [1], [7]],
            [[4], [2], [6], [2]],
            [[4], [2], [1], [2]],
            [[4], [2], [1], [2]],
        ], dtype=tf.float32
    )
    ds = tf.data.Dataset.from_tensor_slices((x, y))
    ds = ds.batch(n_example, drop_remainder=False)

    # A minimum model to reproduce the issue.
    bar_cell = BarCell()
    model = FooModel(cell=bar_cell)
    compile_kwargs = {
        'loss': tf.keras.losses.MeanSquaredError(),
        'optimizer': tf.keras.optimizers.Adam(learning_rate=.001),
    }
    model.compile(**compile_kwargs)

    # Call fit which will trigger gradient computations and raise an error
    # during graph execution.
    model.fit(ds, epochs=1)