Apply LSTM to each matrix element with Keras I'm trying to apply a LSTM/GRU to each entry of a matrix $X$ 

note: Each matrix element is a time-series, so shape of X is (batch_size, rows, cols, time_steps, dims)

$
y_{i,j}=
    \begin{cases}
      0, &amp; \text{if}\ x_{i,j}\small[0\small] = 0 \\
      LSTM(x_{i,j}), &amp; \text{otherwise}
    \end{cases}
$

, where:

$x_{i,j} := \text{element (i,j)  of matrix X. This element is a time series}$
$x_{i,j}\small{[0]} := \text{ First tiem step of time series $x_{i,j}$ }$

$y_{i,j} := \text{LSTM output for $x_{i,j}$}$

PROBLEM

I've managed to evaluate the condition and to apply the LSTM to each entry, but I can not build a keras model with these outputs, beacuse of the typical error:

AttributeError: 'NoneType' object has no attribute '_inbound_nodes' 


CODE

import keras

def convert_to_layer(X):
    return X

def gather(r):
    def gather_(X):
        return X[:,r,:]
    return gather_

n_rows = 2
n_cols = 2
time_steps = 20
dims = 2 
units = 32

M = keras.layers.Input(batch_shape=(1,n_rows,n_cols,time_steps,dims))

# Serialize inputs
M2 = keras.layers.Reshape((n_rows*n_cols,time_steps,dims))(M)

# Gather element 0 (only work on this element for simplicity...)
seq = keras.layers.Lambda(gather(0))(M2)

# Create condition
cond = keras.backend.equal(keras.backend.sum(keras.backend.abs(seq),axis=-1),0)

# Apply if statement
y = keras.backend.switch(cond, 
                          lambda: keras.layers.GRU(units,return_sequences=True)(seq), 
                          lambda: keras.backend.zeros(shape=(1,time_steps,units)))

y = keras.layers.Lambda(convert_to_layer)(y)
print(y)

model = keras.models.Model(M,y)


OUTPUT

Tensor("lambda_311/Identity:0", shape=(1, 20, 32), dtype=float32)

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

&lt;ipython-input-76-983e9545cca1&gt; in &lt;module&gt;()
     29 print(y)
     30 
---&gt; 31 model = keras.models.Model(M,y)

5 frames

/usr/local/lib/python3.6/dist-packages/keras/engine/network.py in build_map(tensor, finished_nodes, nodes_in_progress, layer, node_index, tensor_index)
   1378             ValueError: if a cycle is detected.
   1379         """
-&gt; 1380         node = layer._inbound_nodes[node_index]
   1381 
   1382         # Prevent cycles.

AttributeError: 'NoneType' object has no attribute '_inbound_nodes'


QUESTIONS

The problem is with the keras.backend.switch (if the switch is removed and the output of the model is the output of the gru_cell unconditionally, the model can be build)


Why isn't the model compiling?
Any other alternative way of computing the LSTM over the matrix $X$
SOLUTION

Replace:

# Apply if statement
y = keras.backend.switch(cond, 
                          lambda: keras.layers.GRU(units,return_sequences=True)(seq), 
                          lambda: keras.backend.zeros(shape=(1,time_steps,units)))

y = keras.layers.Lambda(convert_to_layer)(y)


with:

y = keras.layers.Lambda(lambda x: keras.backend.switch(cond, 
                          lambda: keras.layers.GRU(units,return_sequences=True)(x), 
                          lambda: keras.backend.zeros(shape=(1,time_steps,units))))(seq)