Inserting an LSTM layer raises a shape error

[arvoelke]

Steps to reproduce:

• Open docs/examples/lmu.ipynb
• Replace the network definition with the following:

with nengo.Network(seed=seed) as net:
    nengo_dl.configure_settings(
        trainable=None, stateful=False, keep_history=False,
    )
    
    inp = nengo.Node(np.zeros(train_images.shape[-1]))
    h = nengo_dl.Layer(tf.keras.layers.LSTM(units=128))(inp)
    out = nengo_dl.Layer(tf.keras.layers.Dense(units=10))(h)
    p = nengo.Probe(out)

Build finished in 0:00:00                                                      
Optimization finished in 0:00:00                                               
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-25-fc54bb9ca5d2> in <module>
     12 
     13 with nengo_dl.Simulator(
---> 14         net, minibatch_size=100, unroll_simulation=8) as sim:
     15     sim.compile(
     16         loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),

~/git/nengo-dl/nengo_dl/simulator.py in __init__(self, network, dt, seed, model, device, unroll_simulation, minibatch_size, progress_bar)
    510         # build keras models
    511         self.graph = tf.Graph()
--> 512         self._build_keras()
    513 
    514         # initialize sim attributes

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/nengo/utils/magic.py in __call__(self, *args, **kwargs)
    179                 return self.wrapper(wrapped, instance, args, kwargs)
    180             else:
--> 181                 return self.wrapper(self.__wrapped__, self.instance, args, kwargs)
    182         else:
    183             instance = getattr(self.__wrapped__, "__self__", None)

~/git/nengo-dl/nengo_dl/simulator.py in with_self(wrapped, instance, args, kwargs)
     48         instance.tensor_graph.device
     49     ):
---> 50         output = wrapped(*args, **kwargs)
     51     tf.keras.backend.set_floatx(keras_dtype)
     52 

~/git/nengo-dl/nengo_dl/simulator.py in _build_keras(self)
    535                 # if the global learning phase is set, use that
    536                 training=backend._GRAPH_LEARNING_PHASES.get(
--> 537                     backend._DUMMY_EAGER_GRAPH, None
    538                 ),
    539             ),

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py in __call__(self, inputs, *args, **kwargs)
    845                     outputs = base_layer_utils.mark_as_return(outputs, acd)
    846                 else:
--> 847                   outputs = call_fn(cast_inputs, *args, **kwargs)
    848 
    849             except errors.OperatorNotAllowedInGraphError as e:

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/autograph/impl/api.py in wrapper(*args, **kwargs)
    290   def wrapper(*args, **kwargs):
    291     with ag_ctx.ControlStatusCtx(status=ag_ctx.Status.DISABLED):
--> 292       return func(*args, **kwargs)
    293 
    294   if inspect.isfunction(func) or inspect.ismethod(func):

~/git/nengo-dl/nengo_dl/tensor_graph.py in call(self, inputs, training, progress, stateful)
    400         with progress.sub("build stage", max_value=len(self.plan) * self.unroll) as sub:
    401             steps_run, probe_arrays, final_internal_state = (
--> 402                 self._build_loop(sub) if self.use_loop else self._build_no_loop(sub)
    403             )
    404 

~/git/nengo-dl/nengo_dl/tensor_graph.py in _build_loop(self, progress)
    514             loop_vars=loop_vars,
    515             parallel_iterations=1,  # TODO: parallel iterations work in eager mode
--> 516             back_prop=not self.inference_only,
    517         )
    518 

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop_v2(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, maximum_iterations, name)
   2476       name=name,
   2477       maximum_iterations=maximum_iterations,
-> 2478       return_same_structure=True)
   2479 
   2480 

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in while_loop(cond, body, loop_vars, shape_invariants, parallel_iterations, back_prop, swap_memory, name, maximum_iterations, return_same_structure)
   2751       ops.add_to_collection(ops.GraphKeys.WHILE_CONTEXT, loop_context)
   2752     result = loop_context.BuildLoop(cond, body, loop_vars, shape_invariants,
-> 2753                                     return_same_structure)
   2754     if maximum_iterations is not None:
   2755       return result[1]

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in BuildLoop(self, pred, body, loop_vars, shape_invariants, return_same_structure)
   2243       with ops.get_default_graph()._mutation_lock():  # pylint: disable=protected-access
   2244         original_body_result, exit_vars = self._BuildLoop(
-> 2245             pred, body, original_loop_vars, loop_vars, shape_invariants)
   2246     finally:
   2247       self.Exit()

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py in _BuildLoop(self, pred, body, original_loop_vars, loop_vars, shape_invariants)
   2168         expand_composites=True)
   2169     pre_summaries = ops.get_collection(ops.GraphKeys._SUMMARY_COLLECTION)  # pylint: disable=protected-access
-> 2170     body_result = body(*packed_vars_for_body)
   2171     post_summaries = ops.get_collection(ops.GraphKeys._SUMMARY_COLLECTION)  # pylint: disable=protected-access
   2172     if not nest.is_sequence_or_composite(body_result):

~/git/nengo-dl/nengo_dl/tensor_graph.py in loop_body(loop_i, n_steps, probe_arrays, saved_state, base_params)
    486                         )
    487 
--> 488             loop_i = self._build_inner_loop(loop_i, update_probes, progress)
    489 
    490             state_arrays = tuple(self.signals.bases[key] for key in self.saved_state)

~/git/nengo-dl/nengo_dl/tensor_graph.py in _build_inner_loop(self, loop_i, update_probes, progress)
    658                 with tf.control_dependencies([loop_i]):
    659                     # build operators
--> 660                     side_effects = self.op_builder.build(progress)
    661 
    662                     logger.debug("collecting probe tensors")

~/git/nengo-dl/nengo_dl/builder.py in build(self, progress)
     98 
     99             with self.name_scope(ops):
--> 100                 output = self.op_builds[ops].build_step(self.signals)
    101 
    102             if isinstance(output, (tf.Tensor, tf.Variable)):

~/git/nengo-dl/nengo_dl/tensor_node.py in build_step(self, signals)
    352             if len(inputs) == 1:
    353                 inputs = inputs[0]
--> 354             output = self.func.call(inputs)
    355         else:
    356             output = self.func(*inputs)

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/recurrent_v2.py in call(self, inputs, mask, training, initial_state)
    918           input_length=timesteps,
    919           time_major=self.time_major,
--> 920           zero_output_for_mask=self.zero_output_for_mask)
    921       runtime = _runtime(_RUNTIME_UNKNOWN)
    922     else:

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/keras/backend.py in rnn(step_function, inputs, initial_states, go_backwards, mask, constants, unroll, input_length, time_major, zero_output_for_mask)
   3902 
   3903   for input_ in flatted_inputs:
-> 3904     input_.shape.with_rank_at_least(3)
   3905 
   3906   if mask is not None:

~/anaconda3/envs/nengo-dl/lib/python3.7/site-packages/tensorflow_core/python/framework/tensor_shape.py in with_rank_at_least(self, rank)
   1030     """
   1031     if self.rank is not None and self.rank < rank:
-> 1032       raise ValueError("Shape %s must have rank at least %d" % (self, rank))
   1033     else:
   1034       return self

ValueError: Shape (1, 100) must have rank at least 3

I have also tried adding unroll=True to the LSTM and/or configuring stateful=True and/or configuring keep_history=True under nengo_dl.configure_settings.

[drasmuss]

If you use h = nengo_dl.Layer(tf.keras.layers.LSTM(units=128))(inp, shape_in=(n_steps, d)) that should work (where n_steps is the number of timesteps in your data and d is the dimensionality of the data on each timestep).

[arvoelke]

Thanks. Got this to work with the following code:

with nengo.Network(seed=seed) as net:
    nengo_dl.configure_settings(
        trainable=None, stateful=False, keep_history=False,
    )
   
    inp = nengo.Node(np.zeros(np.prod(train_images.shape[1:])))

    h = nengo_dl.Layer(tf.keras.layers.LSTM(units=128))(
        inp, shape_in=(train_images.shape[1], 1))

    out = nengo_dl.Layer(tf.keras.layers.Dense(units=10))(h)
    p = nengo.Probe(out)

and using train_images.reshape((train_images.shape[0], 1, -1)) in place of train_images where passed to sim. Note this passes the entire sequence in one step, rather than the usual Nengo approach of iterating across each time-step in the simulation.

[drasmuss]

Marking this as resolved since the fix above works!