Split channels in 3 parts. Shifts 1st and 3rd sections to left/right.

def DiagonalGate(x, params, **kwargs):
  """Split channels in 3 parts. Shifts 1st and 3rd sections to left/right."""
  del params
  del kwargs
  # x : [batch, 1, length, depth]
  x = np.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 np.concatenate(xs, axis=3)