def MakeZeroState(depth_multiplier=1):
  """Makes zeros of shape like x but removing the length (axis 1)."""
  def f(x):  # pylint: disable=invalid-name
    if len(x.shape) != 3:
      raise ValueError(f'Layer input should be a rank 3 tensor representing'
                       f' (batch_size, sequence_length, feature_depth); '
                       f'instead got shape {x.shape}.')
    return jnp.zeros((x.shape[0], depth_multiplier * x.shape[-1]),
                     dtype=jnp.float32)
  return base.Fn('MakeZeroState', f)
Makes zeros of shape like x but removing the length (axis 1).
trax/layers/rnn.py