LMU allows the input to have ndim>3 (and then fails later)

[arvoelke]

Minimal reproducer:

inp = tf.keras.layers.Input(shape=(1000, 8, 128))

out = keras_lmu.LMU(
    memory_d=128,
    order=6,
    theta=10,
    hidden_cell=tf.keras.layers.Dense(100),
)(inp)

model = tf.keras.Model(inp, out)
model.summary()

Output:

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_0 (InputLayer)         [(None, 1000, 8, 128)]    0         
_________________________________________________________________
lmu_0 (LMU)                  (None, 100)               93326     
=================================================================
Total params: 93,326
Trainable params: 93,284
Non-trainable params: 42

It looks like the extra dimensions are being collapsed somewhere internally through a reshape, but attempting to use the model can result in shaping errors.

If you try this same thing with other RNNs, such as out = tf.keras.layers.SimpleRNN(128)(inp) then you will get the error message:

ValueError: Input 0 of layer simple_rnn_0 is incompatible with the layer: expected ndim=3, found ndim=4. Full shape received: [200, 1000, 8, 128]

I think we should perform the same kind of validation to disallow ndim>3 explicitly (unless of course there's a way to make this work correctly)?