def RepresentationMask(serializer):
  """Upsamples a mask to cover the serialized representation."""
  # Trax enforces the mask to be of the same size as the target. Get rid of the
  # extra dimensions.
  def representation_mask(mask):
    # mask shape (batch_size,4)
    mask = jnp.amax(mask, axis=tuple(range(2, mask.ndim)))
    # mask shape (batch_size,4)
    mask = jnp.repeat(
        mask[..., jnp.newaxis],
        repeats=serializer.representation_length,
        axis=2)
    # mask shape (batch_size,4,representation_length)
    return mask
  return tl.Fn('RepresentationMask', representation_mask)
Upsamples a mask to cover the serialized representation.
trax/rl/serialization_utils.py