LMUFFT cannot be created if hidden_cell is Dense

[arvoelke]

Versions:

• keras_lmu==0.3.2.dev0
• tensorflow==2.4.1

Taking this example from the unit tests:

keras-lmu/keras_lmu/tests/test_layers.py

Lines 158 to 164 in ab07757

	out = layers.LMUFFT(
	1,
	2,
	3,
	tf.keras.layers.SimpleRNNCell(4),
	return_sequences=True,
	)(inp)

and modifying it as follows:

        out = layers.LMUFFT(
            1,
            2,
            3,
            tf.keras.layers.Dense(4),  # tf.keras.layers.SimpleRNNCell(4),
            return_sequences=True,
        )(inp)

results in the error:

ValueError: in user code:

    /home/arvoelke/git/keras-lmu/keras_lmu/layers.py:660 call  *
        h = tf.keras.layers.TimeDistributed(self.hidden_cell)(
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py:1012 __call__  **
        outputs = call_fn(inputs, *args, **kwargs)
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/layers/wrappers.py:244 call
        output_shape = self.compute_output_shape(input_shape).as_list()
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/layers/wrappers.py:188 compute_output_shape
        child_output_shape = self.layer.compute_output_shape(child_input_shape)
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/layers/core.py:1218 compute_output_shape
        raise ValueError(

    ValueError: The innermost dimension of input_shape must be defined, but saw: (None, None)

../../anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py:670: ValueError

[NarsimhaChilkuri]

Although TimeDistributed(Dense(...)) and Dense(...) are equivalent, replacing

keras-lmu/keras_lmu/layers.py

Lines 660 to 662 in ab07757

	h = tf.keras.layers.TimeDistributed(self.hidden_cell)(
	h_in, training=training
	)

with

                h = self.hidden_cell(
                    h_in, training=training
                )

fixes the issue.