code: 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)
  
imports: from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: DiagonalGate() and returns: tl, jnp