def DiagonalGate():
  """Split channels in 3 parts. Shifts 1st and 3rd sections to left/right."""

  def f(x):  # pylint: disable=invalid-name
    # x : [batch, 1, length, depth]
    x = jnp.pad(x, [(0, 0), (0, 0), (1, 1), (0, 0)],
                mode='constant', constant_values=0.0)
    depth = x.shape[-1] // 3
    assert 3 * depth == x.shape[-1], ('Depth must be divisible by 3', depth,
                                      x.shape)
    xs = [
        x[:, :, :-2, :depth], x[:, :, 1:-1, depth:2 * depth],
        x[:, :, 2:, 2 * depth:3 * depth]
    ]
    return jnp.concatenate(xs, axis=3)
  return tl.Fn('DiagonalGate', f)
Split channels in 3 parts. Shifts 1st and 3rd sections to left/right.
trax/models/neural_gpu.py