# Ensure x is a 2D matrix (batch, features)
chex.assert_rank(x, 2)

# Ensure x is either 2D or 3D
chex.assert_rank(x, (2, 3))

        __
        __