def Deinterleave(x_size, y_size):
  """Layer that does the inverse of Interleave."""
  def deinterleave(inputs):
    reprs = inputs
    (batch_size, length) = reprs.shape[:2]
    shape_suffix = reprs.shape[2:]
    remainder_length = length % (x_size + y_size)
    if remainder_length > 0:
      remainder = reprs[:, None, -remainder_length:]
      reprs = reprs[:, :-remainder_length]
    reprs = jnp.reshape(reprs, (batch_size, -1, x_size + y_size) + shape_suffix)
    x_reprs = reprs[:, :, :x_size]
    y_reprs = reprs[:, :, x_size:]
    if remainder_length > 0:
      x_reprs = jnp.concatenate((x_reprs, remainder), axis=1)
    return (x_reprs, y_reprs)
  return tl.Fn('Deinterleave', deinterleave, n_out=2)
Layer that does the inverse of Interleave.
trax/rl/serialization_utils.py