Allow variable layout to be set on construction. (#21264)

If a variable has a particular layout (e.g. is sharded across
multiple devices/hosts), then any corresponding auxiliary variables,
like optimizer gradient accumulators, need to have the same layout.
This requires use to set the variable layout prior to initialization
so that it is initialized correctly and efficiently across devices.

Added optional `kwargs` to the base `Variable` class so they can
handle (and ignore) any backend-specific options.  Modified
`JaxVariable` to allow setting the layout parameter on construction.
Modified `add_variable_from_reference` to copy the layout from
the reference variable.

Author: cantonios

keras/src/backend/common/variables.py

@@ -51,6 +51,7 @@ class Variable:
         value: The current value of the variable (NumPy array or tensor).
         name: The name of the variable (string).
         path: The path of the variable within the Keras model or layer (string).
+        kwargs: Additional backend-specific keyword arguments.
 
     Examples:
 
@@ -98,7 +99,9 @@ def __init__(
         aggregation="none",
         synchronization="auto",
         name=None,
+        **kwargs,
     ):
+        del kwargs
         name = name or auto_name(self.__class__.__name__)
         if not isinstance(name, str) or "/" in name:
             raise ValueError(

keras/src/backend/jax/core.py

@@ -20,20 +20,30 @@
 
 
 class Variable(KerasVariable):
+    def __init__(self, *args, layout=None, **kwargs):
+        # Intercept layout parameter so that it is available
+        # during initialization.
+        self._layout = layout
+        super().__init__(*args, **kwargs)
+
     def _initialize(self, value):
         # Note that variable.shape is needed by distribution_lib
         self._shape = self._validate_shape(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.get_variable_layout(self).backend_layout
-        else:
-            self._layout = None
+        if self._layout is None and distribution is not None:
+            tensor_layout = distribution.get_variable_layout(self)
+            from keras.src.distribution import TensorLayout
+
+            if isinstance(tensor_layout, TensorLayout):
+                self._layout = tensor_layout.backend_layout
+            else:
+                self._layout = tensor_layout
         self._direct_assign(value)
 
     def _direct_assign(self, value):
-        if getattr(self, "_layout", None) is not None:
+        if self._layout is not None:
             value = distribution_lib.distribute_variable(value, self._layout)
         self._value = value
 

keras/src/optimizers/base_optimizer.py

@@ -244,6 +244,7 @@ def add_variable(
         initializer="zeros",
         dtype=None,
         aggregation="none",
+        layout=None,
         name=None,
     ):
         """Add a variable to the optimizer.
@@ -261,6 +262,7 @@ def add_variable(
                 the type of multi-replica aggregation to be used for this
                 variable when writing custom data parallel training loops.
                 Defaults to `"none"`.
+            layout: Optional tensor layout.  Defaults to `None`.
             name: String name of the variable. Useful for debugging purposes.
 
         Returns:
@@ -275,6 +277,7 @@ def add_variable(
                 dtype=dtype,
                 trainable=False,
                 aggregation=aggregation,
+                layout=layout,
                 name=name,
             )
         self._track_variable(variable)
@@ -319,6 +322,7 @@ def add_variable_from_reference(
             initializer=initializer,
             dtype=reference_variable.dtype,
             name=name,
+            layout=getattr(reference_variable, "_layout", None),
         )
 
     def add_optimizer_variables(