Apply LSTM to each matrix element with Keras

Question

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, & \text{if}\ x_{i,j}\small[0\small] = 0 \\ LSTM(x_{i,j}), & \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)

<ipython-input-76-983e9545cca1> in <module>()
     29 print(y)
     30 
---> 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         """
-> 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)

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

Answer 1

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)