code: def DiagonalGate(x, params, **kwargs):
    
    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)
docs: Split channels in 3 parts. Shifts 1st and 3rd sections to left/right.