将Lambda函数重写为自定义keras层

Question

要在TfLite中使用LSTM，我需要将Lambda层重写为tf keras自定义Layer对象，以便它可以容纳权重（因为Lambda层无法做到这一点）。我不太确定如何使用此处的示例。谁能帮助我，并向我展示如何正确地重写为custom layer？

def buildLstmLayer(inputs):

  lstm_cells = []

  lstm_cells.append(TFLiteLSTMCell(256, forget_bias=0, name='rnn0'))
  lstm_cells.append(TFLiteLSTMCell(128, forget_bias=0, name='rnn1'))
  lstm_layers = tf.keras.layers.StackedRNNCells(lstm_cells)

  outputs, _ = dynamic_rnn(
      lstm_layers,
      inputs,
      dtype='float32',
      time_major=True)

  return outputs