keras/src/backend/jax/core.py/Variable/_direct_assign
class Variable:    def _direct_assign(self, value):
        if getattr(self, "_layout", None) is not None:
            value = distribution_lib.distribute_variable(value, self._layout)
        self._value = value