# Ensure axis 1 has exactly 12 elements (e.g., 12 financial features)
chex.assert_axis_dimension(x, axis=1, expected=12)

        __
        __