code: def test_masking(self):
      class MaskedDense(layers.Wrapper):
          def __init__(self, units, **kwargs):
              layer = layers.Dense(
                  units,
                  kernel_initializer=initializers.Constant(0.01),
                  use_bias=False,
              )
              super().__init__(layer, **kwargs)
              self.supports_masking = True
  
          def call(self, inputs, training=False, mask=None):
              unmasked = self.layer.call(inputs)
              if mask is None:
                  return unmasked
              else:
                  return ops.transpose(
                      ops.transpose(unmasked)
                      * ops.cast(mask, inputs.dtype)
                  )
  
      sequence = np.arange(24).reshape((3, 2, 4)).astype("float32")
      layer = layers.TimeDistributed(layer=MaskedDense(1))
      mask = np.array([[False, True], [True, False], [True, True]])
      output = layer(sequence, mask=mask)
      self.assertAllClose(
          np.array([[[0], [0.22]], [[0.38], [0]], [[0.7], [0.86]]]),
          output,
      )
  
file_path: keras/keras/layers/rnn/time_distributed_test.py
imports: import numpy
  import pytest
  from keras import backend
  from keras import initializers
  from keras import layers
  from keras import ops
  from keras import testing
label_desc: Use image node_id 4 for calling the TimeDistributedTest obj's underlying member method code with example usage: obj.test_masking() and returns: unmasked, ops