image/svg+xml Axis(0) Axis(1) Axis(2) .split_at(Axis(2), 2) Input shape: (3, 5, 5)Output shapes: (3, 5, 2) and (3, 5, 3)