# Ensure the array is float32, not int32 or float64
chex.assert_type(x, jnp.float32)

        __
        __