keras-nightly 3.12.0.dev2025092403__py3-none-any.whl → 3.14.0.dev2026010104__py3-none-any.whl

keras/src/backend/common/name_scope.py

@@ -58,7 +58,8 @@ class name_scope:
             name_scope_stack = global_state.get_global_attribute(
                 "name_scope_stack"
             )
-            name_scope_stack.pop()
+            if name_scope_stack:
+                name_scope_stack.pop()
 
 
 def current_path():

keras/src/backend/common/variables.py

@@ -276,13 +276,13 @@ class Variable:
         return self._maybe_autocast(self._value)
 
     def assign(self, value):
-        value = self._convert_to_tensor(value, dtype=self.dtype)
+        value = self._convert_to_tensor(value, dtype=self._dtype)
         if not shape_equal(value.shape, self.shape):
             raise ValueError(
                 "The shape of the target variable and "
                 "the shape of the target value in "
                 "`variable.assign(value)` must match. "
-                f"variable.shape={self.value.shape}, "
+                f"variable.shape={self.shape}, "
                 f"Received: value.shape={value.shape}. "
                 f"Target variable: {self}"
             )
@@ -399,7 +399,11 @@ class Variable:
     def __repr__(self):
         value = None
         if hasattr(self, "_value") and self._value is not None:
-            value = backend.core.convert_to_numpy(self._value)
+            try:
+                value = backend.core.convert_to_numpy(self._value)
+            except:
+                # In some cases the conversion to numpy can fail.
+                pass
         value_str = f", value={value}" if value is not None else ""
         return (
             f"<Variable path={self.path}, shape={self.shape}, "
@@ -595,33 +599,46 @@ def standardize_shape(shape):
                 # `tf.TensorShape` may contain `Dimension` objects.
                 # We need to convert the items in it to either int or `None`
                 shape = shape.as_list()
-        shape = tuple(shape)
 
     if config.backend() == "jax":
         # Replace `_DimExpr` (dimension expression) with None
+        from jax import export as jax_export
+
         shape = tuple(
-            [None if "_DimExpr" in str(type(d)) else d for d in shape]
+            None if jax_export.is_symbolic_dim(d) else d for d in shape
         )
 
-    if config.backend() == "torch":
-        # `shape` might be `torch.Size`. We need to convert the items in it to
-        # either int or `None`
-        shape = tuple(map(lambda x: int(x) if x is not None else None, shape))
-
-    for e in shape:
-        if e is None:
+    # Handle dimensions that are not ints and not None, verify they're >= 0.
+    standardized_shape = []
+    for d in shape:
+        if d is None:
+            standardized_shape.append(d)
             continue
-        if not is_int_dtype(type(e)):
+
+        # Reject these even if they can be cast to int successfully.
+        if isinstance(d, (str, float)):
             raise ValueError(
                 f"Cannot convert '{shape}' to a shape. "
-                f"Found invalid entry '{e}' of type '{type(e)}'. "
+                f"Found invalid dimension '{d}' of type '{type(d)}'. "
             )
-        if e < 0:
+
+        try:
+            # Cast numpy scalars, tf constant tensors, etc.
+            d = int(d)
+        except Exception as e:
+            raise ValueError(
+                f"Cannot convert '{shape}' to a shape. "
+                f"Found invalid dimension '{d}' of type '{type(d)}'. "
+            ) from e
+        if d < 0:
             raise ValueError(
                 f"Cannot convert '{shape}' to a shape. "
                 "Negative dimensions are not allowed."
             )
-    return shape
+        standardized_shape.append(d)
+
+    # This also turns subclasses of `tuple` (e.g. `torch.Size`) to plain tuple.
+    return tuple(standardized_shape)
 
 
 def shape_equal(a_shape, b_shape):

keras/src/backend/jax/core.py

@@ -30,9 +30,7 @@ class JaxVariable(KerasVariable):
         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)
+    def _initialize_layout(self):
         # We can't import the keras/distribution/distribution_lib
         # due to circular dependency.
         distribution = global_state.get_global_attribute("distribution")
@@ -44,8 +42,28 @@ class JaxVariable(KerasVariable):
                 self._layout = tensor_layout.backend_layout
             else:
                 self._layout = tensor_layout
+
+    def _initialize(self, value):
+        # Note that variable.shape is needed by distribution_lib
+        self._shape = self._validate_shape(value.shape)
+        self._initialize_layout()
         self._direct_assign(value)
 
+    def _initialize_with_initializer(self, initializer):
+        self._initialize_layout()
+        layout = self._layout
+        shape = self._shape
+        if should_shard_at_init(layout, shape):
+            jitted_initializer = jax.jit(
+                initializer.__call__,
+                out_shardings=layout,
+                static_argnames=["shape", "dtype"],
+            )
+            value = jitted_initializer(shape=self._shape, dtype=self._dtype)
+            self._value = value
+        else:
+            super()._initialize_with_initializer(initializer)
+
     def _direct_assign(self, value):
         if self._layout is not None:
             value = distribution_lib.distribute_variable(value, self._layout)
@@ -112,6 +130,12 @@ if config.is_nnx_enabled():
             # The real value is now set in self._value, sync it to raw_value
             object.__setattr__(self, "raw_value", self._value)
 
+        def _initialize_with_initializer(self, initializer):
+            value = self._convert_to_tensor(
+                initializer(self._shape, dtype=self._dtype)
+            )
+            self._initialize(value)
+
         @property
         def _value(self):
             if hasattr(self, "raw_value"):
@@ -234,6 +258,71 @@ if config.is_nnx_enabled():
 
     Variable = NnxVariable
 
+    def _flatten_nnx_variable(variable):
+        children = (variable.raw_value,)
+        # We copy __dict__ to avoid side effects
+        keras_state = variable.__dict__.copy()
+        # Remove elements that might be problematic or redundant if
+        # nnx.Variable's __getstate__
+        keras_state.pop("raw_value", None)
+        aux_data = (
+            variable._var_metadata,
+            getattr(variable, "_trace_state", None),
+            keras_state,
+        )
+        return children, aux_data
+
+    def _unflatten_nnx_variable(aux_data, children):
+        var_metadata, trace_state, keras_state = aux_data
+        raw_value = children[0]
+
+        # Create uninitialized instance
+        variable = NnxVariable.__new__(NnxVariable)
+
+        # Restore state
+        variable._var_metadata = var_metadata
+        if trace_state is not None:
+            variable._trace_state = trace_state
+        variable.__dict__.update(keras_state)
+        variable.raw_value = raw_value
+
+        return variable
+
+    try:
+        jax.tree_util.register_pytree_node(
+            NnxVariable,
+            _flatten_nnx_variable,
+            _unflatten_nnx_variable,
+        )
+    except ValueError:
+        pass
+
+    def __setattr__(self, name, value):
+        # Mirror Keras attributes to _var_metadata to ensure persistence
+        # if the Pytree registration is not respected by NNX.
+        if (
+            name != "_var_metadata"
+            and name not in ("_raw_value", "_trace_state")
+            and hasattr(self, "_var_metadata")
+        ):
+            self._var_metadata[name] = value
+
+        object.__setattr__(self, name, value)
+
+    NnxVariable.__setattr__ = __setattr__
+
+
+def should_shard_at_init(init_layout, shape):
+    if not isinstance(init_layout, jax.sharding.NamedSharding):
+        return False
+
+    if all(dim is None for dim in init_layout.spec):
+        return False
+
+    size_threshold = 250 * 1024 * 1024
+    array_size = np.prod(shape) * 4
+    return array_size >= size_threshold
+
 
 def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
     if ragged:

keras/src/backend/jax/distribution_lib.py

@@ -27,6 +27,20 @@ def list_devices(device_type=None):
     return [f"{device.platform}:{device.id}" for device in jax_devices]
 
 
+def get_device_count(device_type=None):
+    """Returns the number of available JAX devices.
+    Args:
+        device_type: Optional device type to count (e.g., "cpu", "gpu", "tpu").
+            If `None`, it defaults to counting "gpu" or "tpu" devices if
+            available, otherwise it counts "cpu" devices. It does not
+            return the sum of all device types.
+    Returns:
+        int: The total number of JAX devices for the specified type.
+    """
+    device_type = device_type.lower() if device_type else None
+    return jax.device_count(device_type)
+
+
 def distribute_variable(value, layout):
     """Create a distributed variable for JAX.
 
@@ -146,13 +160,13 @@ def initialize_rng():
     # Check if the global seed generator is set and ensure it has an initialized
     # seed.  Otherwise, reset the seed to the global seed.
     global_seed_generator = global_state.get_global_attribute(
-        "global_seed_generator"
+        seed_generator.GLOBAL_SEED_GENERATOR
     )
     if global_seed_generator is not None:
         seed = global_seed_generator.get_config()["seed"]
         if seed is None:
             global_state.set_global_attribute(
-                "global_seed_generator",
+                seed_generator.GLOBAL_SEED_GENERATOR,
                 seed_generator.SeedGenerator(
                     seed=global_seed,
                     name=global_seed_generator.name,

keras/src/backend/jax/linalg.py

@@ -97,3 +97,7 @@ def lstsq(a, b, rcond=None):
     a = convert_to_tensor(a)
     b = convert_to_tensor(b)
     return jnp.linalg.lstsq(a, b, rcond=rcond)[0]
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    return jax.jvp(fun, primals, tangents, has_aux=has_aux)