Support custom cell/RNN layers with extension types

Issue type

Feature Request

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.15

Custom code

Yes

OS platform and distribution

Windows 11

Mobile device

No response

Python version

3.11

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

I want to write a Keras-like model with keras.layers.RNN that supports Extension types, both for inputs and states.

Standalone code to reproduce the issue

import keras
import tensorflow as tf

class MaskedTensor(tf.experimental.ExtensionType):
    """A tensor paired with a boolean mask, indicating which values are valid."""

    values: tf.Tensor
    mask: tf.Tensor
    shape: tf.TensorShape
    dtype: tf.DType

    def __init__(self, values, mask):
        self.values = values
        self.mask = mask
        self.shape = values.shape
        self.dtype = values.dtype

@tf.experimental.dispatch_for_api(tf.compat.v1.transpose)
def transpose(a: MaskedTensor, perm=None, name="transpose", conjugate=False):
    values = tf.transpose(a.values, perm, conjugate, name)
    mask = tf.transpose(a.mask, perm, conjugate, name)
    return MaskedTensor(values, mask)

@tf.experimental.dispatch_for_api(tf.shape)
def shape(input: MaskedTensor, out_type=None, name=None):
    return tf.shape(input.values, out_type, name)

@tf.experimental.dispatch_for_api(tf.unstack)
def unstack(value: MaskedTensor, num=None, axis=0, name="unstack"):
    values = tf.unstack(value.values, num, axis, name)
    mask = tf.unstack(value.mask, num, axis, name)
    return [MaskedTensor(x, m) for x, m in zip(values, mask)]

@keras.saving.register_keras_serializable()
class Cell(tf.keras.layers.Layer):
    @property
    def state_size(self):
        return tf.TensorShape([5])

    def call(self, inputs, states):
        assert isinstance(inputs, MaskedTensor)
        assert isinstance(states, MaskedTensor)
        return inputs, states

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        return MaskedTensor(tf.zeros((batch_size, 5)), tf.ones((batch_size, 5), tf.bool))

if __name__ == "__main__":
    input_spec = MaskedTensor.Spec(
        values=tf.TensorSpec(shape=[2, 10, 5]),
        mask=tf.TensorSpec(shape=[2, 10, 5]),
        shape=[2, 10, 5],
        dtype=tf.float32,
    )
    x = tf.keras.layers.Input(type_spec=input_spec)
    y = tf.keras.layers.RNN(Cell(), return_sequences=True, stateful=True, unroll=True)(x)
    model = tf.keras.models.Model(x, y)
    model.summary()

Relevant log output

File "D:\projects\testing\proof_of_concept.py", line 62, in <module>
    y = tf.keras.layers.RNN(Cell(), return_sequences=True, stateful=True, unroll=True)(x)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\.envs\tensorflow\Lib\site-packages\keras\src\layers\rnn\base_rnn.py", line 557, in __call__
    return super().__call__(inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\.envs\tensorflow\Lib\site-packages\keras\src\utils\traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "D:\.envs\tensorflow\Lib\site-packages\tensorflow\python\framework\constant_op.py", line 103, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Attempt to convert a value (MaskedTensor(values=<tf.Tensor: shape=(2, 5), dtype=float32, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float32)>, mask=<tf.Tensor: shape=(2, 5), dtype=bool, numpy=
array([[ True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True]])>, shape=TensorShape([2, 5]), dtype=tf.float32)) with an unsupported type (<class '__main__.MaskedTensor'>) to a Tensor.

Hi @Johansmm -

Thanks for reporting the issue. Here in your code there are multiple corrections required.

If you want to create MaskedTensor using ExtensionType, then you can can use it like this:

class MaskedTensor(tf.experimental.ExtensionType):
"""A tensor paired with a boolean mask, indicating which values are valid."""

values: tf.Tensor
mask: tf.Tensor

After defining MaskedTensor with above need to remove shape and dtype argument from input_spec
At this line x = tf.keras.layers.Input(type_spec=input_spec), Input layer doesn't have type_spec argument. So need to correct like this: x = tf.keras.layers.Input(batch_size=2,shape=(10,5)).

Attached gist for your reference.

Hi @mehtamansi29 thanks for your response. However, I don't think I have been clear in explaining what the purpose of the code is: I want to know how I can write RNN models with custom cell that works with Extension types tensors.

For this, I have written the example to show my problem based on :

How to write Keras models using Extension types: Here, is explained that Keras supports tf.keras.layer.Input with Extension type through type_spec field (as I made in the example). Note in the guide is marked : 'They must have a field or property named shape. shape[0] is assumed to be the batch dimension.'
How to work with RNN layers with custom cells: section define a custom cell that supports nested input/output.

I hope that my model can make inferences with MaskedTensor inputs, but with the corrections you make the following code does not work:

import tensorflow as tf
mt = MaskedTensor(tf.random.uniform((2,10,5)), tf.ones((2,10,5)))
model(mt)

Error:

ValueError: Exception encountered when calling Functional.call().

Attempt to convert a value (MaskedTensor(values=<tf.Tensor: shape=(2, 10, 5), dtype=float32, numpy=
array([[[0.28294933, 0.89941823, 0.6277088 , 0.3004167 , 0.4065286 ],
        [0.59596014, 0.7929536 , 0.6331537 , 0.15866613, 0.29780304],
        [0.31396973, 0.87872875, 0.13612425, 0.7689322 , 0.6524085 ],
        [0.65356696, 0.74440706, 0.381616  , 0.3481027 , 0.44483578],
        [0.2988621 , 0.0631609 , 0.5148549 , 0.89417315, 0.4907987 ],
        [0.23625493, 0.78680897, 0.6701437 , 0.2609341 , 0.16422784],
        [0.09855855, 0.5578295 , 0.8797028 , 0.17377365, 0.9087467 ],
        [0.91141987, 0.7385657 , 0.5835092 , 0.5579853 , 0.8384237 ],
        [0.08012462, 0.56617975, 0.700922  , 0.18580115, 0.61618245],
        [0.47631788, 0.3428775 , 0.1811893 , 0.2038809 , 0.19900978]],

       [[0.4667045 , 0.64295805, 0.35533714, 0.46243107, 0.28277063],
        [0.3471583 , 0.32578683, 0.39409876, 0.5108174 , 0.3006178 ],
        [0.5338242 , 0.23265481, 0.06676841, 0.6289011 , 0.10211515],
        [0.9826255 , 0.50816226, 0.995906  , 0.28830194, 0.7350259 ],
        [0.20371187, 0.75276816, 0.03341246, 0.22956371, 0.14091146],
        [0.5101521 , 0.5355145 , 0.77825236, 0.11842644, 0.09967971],
        [0.5528343 , 0.12923944, 0.9135002 , 0.31218648, 0.09520006],
        [0.6237818 , 0.46556568, 0.45628858, 0.22421765, 0.6033968 ],
        [0.10589111, 0.08551514, 0.20975125, 0.5542921 , 0.14889371],
        [0.04052258, 0.3114897 , 0.3219484 , 0.05069757, 0.7502247 ]]],
      dtype=float32)>, mask=<tf.Tensor: shape=(2, 10, 5), dtype=float32, numpy=
array([[[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]],

       [[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]]], dtype=float32)>, shape=TensorShape([2, 10, 5]))) with an unsupported type (<class '__main__.MaskedTensor'>) to a Tensor.

I hope that the purpose of my question is clearer.

Hi @Johansmm -

Here is the reference where you can create custom RNN layer using subclassing.

And for creating RNN models with custom cell that works with Extension types tensors.

class MaskedTensor(tf.experimental.ExtensionType):
  """A tensor paired with a boolean mask, indicating which values are valid."""
  values: tf.Tensor
  mask: tf.Tensor  

@keras.saving.register_keras_serializable()
class RNNCell(keras.layers.Layer):
  def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.state_size = units

  def build(self, input_shape):
      self.kernel = self.add_weight(shape=(input_shape[-1], self.units),
                                    initializer='uniform',
                                    name='kernel')
      self.recurrent_kernel = self.add_weight(
          shape=(self.units, self.units),
          initializer='uniform',
          name='recurrent_kernel')
      self.built = True

  def call(self, inputs, states):
      prev_output = states[0]
      h = ops.matmul(inputs, self.kernel)
      output = h + ops.matmul(prev_output, self.recurrent_kernel)
      return output, [output]

batch_size= 32
input_shape = (10,)
input= keras.Input(shape=input_shape,batch_size=batch_size)

values= tf.random.normal(shape=(batch_size, 10))
mask= tf.random.uniform(shape=(batch_size, 10)) > 0.5
input_spec= MaskedTensor(values=values, mask=mask)

print(type(input_spec))
rnn= keras.layers.RNN(RNNCell(units=32))(input)
model = keras.models.Model(input,rnn)
model.summary()

Attached gist for the reference.

Hi @mehtamansi29, I do not thing to understand your example, because you are creating a RNNCell layer that does not use MaskedTensor in call(). I think the two topics are being separated, but my goal is to write an RNNCell layer whose inputs/states are MaskedTensor.

With your example it is not possible to make an inference with _inputspec:

y = model(input_spec)
# Raise the following error:
# ValueError: Inputs to a layer should be tensors. Got 'MaskedTensor ....

keras-team / keras