Update API

[drasmuss]

Lots of changes to the LMU API to make it more user-friendly, and add support for new features (like multi-dimensional memories and dropout).

Some of the most significant changes are:

• Removed a bunch of elements that we haven't really found to be that useful in our experimentation with LMUs, like individual trainable/initializer arguments for each connection within the LMU, or the LMUCellODE class (although we have plans to implement an improved method for optimizing the memory component in the future)
• Added the ability to selectively enable/disable connections within the LMU between the hidden and memory components. Note that these default to disabled, so the new default LMU will be a trimmed-down version of the old defaults.
• Added support for multi-dimensional memories.
• Added support for arbitrary hidden components (anything that implements the Keras RNNCell API can be used, e.g. tf.keras.layers.LSTMCell).

See the diff in CHANGES.rst for all the details.

[gsmalik]

I had just one small comment about a potential dimension mismatch if hidden_cell is None and input_to_hidden is enabled.

[gsmalik]

In case the hidden_cell is None, then the input_to_hidden connection should always be False maybe? I could imagine a scenario where hidden cell could be None and input_to_hidden could be True without an error being thrown up.

This would set my hidden_output_size to be self.memory_d * self.order (L104) . This would in-turn set my output_size to be self.memory_d * self.order (L126) whereas it should be self.memory_d * self.order + inp_size (L229)

In the case that my hidden_to_memory is enabled then my enc_d would also be set wrong (L143). This I expect would cause dimension mismatch with kernel and actual input

[drasmuss]

I think there are valid reasons to have input_to_hidden with hidden_cell=None (e.g., you could use it as a way to do skip connections in an LMU stack), so I wouldn't want to disable it.

You're definitely right about the dimensions being wrong though, we need to add inp_size as you say. Could you do that, and also add a test for that case (since we're obviously missing that currently!).

[gsmalik]

Yeah, I can try and do this. 👍

[gsmalik]

@drasmuss, I have added a fix in the build part of LMUCell by updating the hidden_output_size and output_size if hidden_cell is None and input_to_hidden connection is enabled.

For the tests, I have added a test that checks the dimensions of the output, hidden_outout_size and output_size when hidden cell is set (RNN/Dense) and not set with inputt_to_hidden connections enabled.

Could you check quickly if this makes sense? I will then merge 👍

[drasmuss]

Pushed a minor change, where we set the output size to None before the layer is built (to make sure that anyone checking that value before building is aware that that value is not correct). And then just some small simplifications to the test. Looks good!