dtype of RNN cell's state is changed to tf.float32 during reset_states

[TillHa]

• Have I written custom code (as opposed to using a stock example script
  provided in Keras): yes
• OS Platform and Distribution: both win10 and CentOS Linux
• TensorFlow installed from: pip
• TensorFlow version: 2.6.0
• Python version: 3.8
• Exact command to reproduce: tf.keras.layers.RNN(cell)

Describe the problem.

I have implemented a recurrent cell which is to be wrapped within a tf.keras.layers.RNN. The cell has a state whose data type is not tf.float32 but tf.complex64. However, each time when layer.reset_states() is invoked, the data type of the state is changed to tf.float32. As a result, a value error is thrown during the initial symbolic call. See attached stack trace.

Describe the current behavior.
The programm crashes at the construction of the RNN layer. See attached stack trace.

Contributing.

I assume, a reason for this issue is line 933, 934 in function reset_states in class RNN in file keras/layers/recurrent.py

      flat_states_variables = tf.nest.map_structure(
          backend.variable, flat_init_state_values)

Here, the initialized state values are stored in flat_init_state_values and backend.variable is called on each of the states. However, no dtype argument is passed to backend.variable. As a consequence it defaults to tf.float32 for all states. T
I would recommend the following patch, which solves the issue for me

      flat_states_variables = tf.nest.map_structure(
    lambda var: backend.variable(var, var.dtype), flat_init_state_values)

I also tried to run the example after replacing the files affected by the latest commit regarding mixed precision. Unfortunately it did not solve the issue for me

Standalone code to reproduce the issue.

Currently, the example fails at the construction of the RNN layer.

import tensorflow as tf

class RecurrentCell(tf.keras.layers.Layer):
    def __init__(self, state_size):
        super(RecurrentCell, self).__init__()
        self.state_size = state_size

    def build(self, input_shape):
        super(RecurrentCell, self).build(input_shape)

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        # explicit initialization with tf.complex64
        return tf.zeros((self.state_size, ), dtype=tf.complex64)

    @tf.function
    def call(self, inputs, states):
        # toy example
        x = inputs
        xfd = tf.signal.rfft(x)[..., :self.state_size]
        yfd = tf.multiply(xfd, states)
        return tf.signal.irfft(yfd), states


recCell = RecurrentCell(state_size=5)

inp = tf.keras.Input(shape=(None, 8),
                     batch_size=32)
out = tf.keras.layers.RNN(recCell,  # crashes
                          return_sequences=True,
                          stateful=True,
                          return_state=False)(inp)
model = tf.keras.Model(inputs=[inp], outputs=[out])

y = model.predict(tf.random.normal((32, 16, 8)))

Source code / Logs
Stacktrace:
stacktrace.txt

[tilakrayal]

@TillHa,
I tried to execute the mentioned code on tf-nightly(2.15.0-dev20230922) without any issue/error. Kindly find the gist of it here. Thank you!

[github-actions]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

[github-actions]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.