keras/src/backend/jax/core.py/Variable/_initialize
class Variable:    def _initialize(self, value):
        value = jnp.array(value, dtype=self._dtype)
        # Note that variable.shape is needed by distribution_lib
        self._shape = tuple(value.shape)
        # We can't import the keras/distribution/distribution_lib
        # due to circular dependency.
        distribution = global_state.get_global_attribute("distribution")
        if distribution is not None:
            self._layout = distribution_lib._to_jax_layout(
                distribution.get_variable_layout(self)
            )
        else:
            self._layout = None
        self._direct_assign(value)