Further changes pertaining to Seed Generator

Author: praveenhosdrug123

keras/src/backend/jax/core.py

@@ -33,41 +33,40 @@ def __init__(self, *args, layout=None, **kwargs):
     def set_tensor_layout(self):
         # We can't import the keras/distribution/distribution_lib
         # due to circular dependency.
-        if self._layout is None:
-            distribution = global_state.get_global_attribute("distribution")
-            if 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
+        distribution = global_state.get_global_attribute("distribution")
+        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
 
     def _initialize(self, value):
         # Note that variable.shape is needed by distribution_lib
         self._shape = self._validate_shape(value.shape)
         self.set_tensor_layout()
         self._direct_assign(value)
 
-    def check_distributed_init(self, initializer):
+    def check_distributed_init(self, initializer, init_layout):
         # Check if 'layout' parameter is supported in the initializer call
         import inspect
 
         sig = inspect.signature(initializer.__call__)
         layout_supported = "layout" in sig.parameters
         # Check if PartitionSpec has any non-None values
-        spec = getattr(self._layout, "spec", None)
+        spec = getattr(init_layout, "spec", None)
         partition_spec = spec if spec is not None else ()
         is_partitioned = any(dim is not None for dim in partition_spec)
-        return layout_supported and is_partitioned
+        return layout_supported and init_layout is not None and is_partitioned
 
     def _initialize_with_initializer(self, initializer):
-        self.set_tensor_layout()
+        init_layout = get_initialization_layout(self.path)
         # Use layout-aware initialization for distributed embeddings
-        if self.check_distributed_init(initializer):
+        if self.check_distributed_init(initializer, init_layout):
             value = self._convert_to_tensor(
-                initializer(self._shape, dtype=self._dtype, layout=self._layout)
+                initializer(self._shape, dtype=self._dtype, layout=init_layout)
             )
         else:
             value = self._convert_to_tensor(
@@ -141,6 +140,12 @@ def __init__(
             # 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"):
@@ -264,6 +269,25 @@ def value(self):
     Variable = NnxVariable
 
 
+def get_initialization_layout(path):
+    distribution = global_state.get_global_attribute("distribution")
+    if distribution is None:
+        return None
+    layout_map = getattr(distribution, "_layout_map", None)
+    if layout_map is None:
+        return None
+    layout_obj = layout_map.get(path)
+    if layout_obj is None:
+        return None
+    from keras.src.distribution import TensorLayout
+
+    if isinstance(layout_obj, TensorLayout):
+        layout_obj = layout_obj.backend_layout
+    if isinstance(layout_obj, jax.sharding.NamedSharding):
+        return layout_obj
+    return None
+
+
 def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
     if ragged:
         raise ValueError("`ragged=True` is not supported with jax backend")

keras/src/backend/jax/distribution_lib.py

@@ -38,10 +38,10 @@ def _distribute_initializer(
     from functools import partial
 
     # Draw seed from the seed generator if seed is not a Jax Array
+    # It is imperative for seed generation to happen before jit compilation
     if seed is None or not isinstance(seed, jax.Array):
-        jax_compatible_seed = seed_generator.draw_seed(None)
-        # Convert to JAX PRNG key format (swap counter and seed value)
-        seed = jax_compatible_seed[::-1]
+        seed = seed_generator.draw_seed(None)[0]
+        seed = jax.random.key(seed)
 
     # Validate all required arguments
     if init_func is None or init_func.func.__name__ not in [

keras/src/backend/jax/random.py

@@ -9,12 +9,14 @@
 
 
 def jax_draw_seed(seed):
-    # Convert to JAX PRNG key format (swap counter and seed value)
     if isinstance(seed, jax.Array):
-        return seed[::-1]
+        if seed.ndim == 0:
+            return jax.random.key(seed)
+        elif seed.ndim == 1 and seed.shape == (2,):
+            return seed
     else:
-        seed_array = draw_seed(seed)
-        return seed_array[::-1]
+        seed = draw_seed(seed)
+        return seed
 
 
 def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None, layout=None):
@@ -40,6 +42,35 @@ def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None, layout=None):
         return sample * stddev + mean
 
 
+def truncated_normal(
+    shape, mean=0.0, stddev=1.0, dtype=None, seed=None, layout=None
+):
+    dtype = dtype or floatx()
+    seed = jax_draw_seed(seed)
+    if layout is not None:
+        from keras.src.backend import distribution_lib
+
+        init_func = partial(
+            jax.random.truncated_normal,
+            shape=shape,
+            dtype=dtype,
+            lower=-2.0,
+            upper=2.0,
+        )
+        return distribution_lib._distribute_initializer(
+            init_func=init_func,
+            mean=mean,
+            stddev=stddev,
+            seed=seed,
+            layout=layout,
+        )
+    else:
+        sample = jax.random.truncated_normal(
+            seed, shape=shape, lower=-2.0, upper=2.0, dtype=dtype
+        )
+        return sample * stddev + mean
+
+
 def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None, layout=None):
     dtype = dtype or floatx()
     seed = jax_draw_seed(seed)
@@ -84,35 +115,6 @@ def randint(shape, minval, maxval, dtype="int32", seed=None):
     )
 
 
-def truncated_normal(
-    shape, mean=0.0, stddev=1.0, dtype=None, seed=None, layout=None
-):
-    dtype = dtype or floatx()
-    seed = jax_draw_seed(seed)
-    if layout is not None:
-        from keras.src.backend import distribution_lib
-
-        init_func = partial(
-            jax.random.truncated_normal,
-            shape=shape,
-            dtype=dtype,
-            lower=-2.0,
-            upper=2.0,
-        )
-        return distribution_lib._distribute_initializer(
-            init_func=init_func,
-            mean=mean,
-            stddev=stddev,
-            seed=seed,
-            layout=layout,
-        )
-    else:
-        sample = jax.random.truncated_normal(
-            seed, shape=shape, lower=-2.0, upper=2.0, dtype=dtype
-        )
-        return sample * stddev + mean
-
-
 def _get_concrete_noise_shape(inputs, noise_shape):
     if noise_shape is None:
         return inputs.shape

keras/src/initializers/random_initializers.py

@@ -6,18 +6,19 @@
 from keras.src.backend import random
 from keras.src.initializers.initializer import Initializer
 from keras.src.saving import serialization_lib
+from keras.src.utils import jax_utils
 
 
 class RandomInitializer(Initializer):
     def __init__(self, seed=None):
         self._init_seed = seed
         if seed is None and backend() == "jax":
-            seed = int(random.draw_seed(None)[0])
+            seed = jax_utils.get_jax_random_seed(seed)
         elif seed is None:
             seed = random.make_default_seed()
         elif isinstance(seed, dict):
             seed = serialization_lib.deserialize_keras_object(seed)
-        elif not isinstance(seed, (int, random.SeedGenerator)):
+        elif not isinstance(seed, (random.SeedGenerator, int)):
             raise ValueError(
                 "`seed` argument should be an instance of "
                 "`keras.random.SeedGenerator()` or an integer. "

keras/src/random/seed_generator.py

@@ -118,29 +118,14 @@ def from_config(cls, config):
 
 def global_seed_generator():
     if jax_utils.is_in_jax_tracing_scope():
-        raise ValueError(
-            "[JAX RNG] When tracing a JAX function, "
-            "you should only use seeded random ops, e.g. "
-            "you should create a `SeedGenerator` instance, attach it "
-            "to your layer/model, and pass the instance as the `seed` "
-            "argument when calling random ops. Unseeded random ops "
-            "would get incorrectly traced by JAX and would become constant "
-            "after tracing. Example:\n\n"
-            "```\n"
-            "# Make sure to set the seed generator as a layer attribute\n"
-            "self.seed_generator = keras.random.SeedGenerator(seed=1337)\n"
-            "...\n"
-            "out = keras.random.normal(shape=(1,), seed=self.seed_generator)\n"
-            "```"
-        )
+        # When we are in Jax Tracing mode, we provide a lightweight
+        # object of the shape and dtype expected
+        return jax_utils.JAXTracingSeedGenerator()
+
     gen = global_state.get_global_attribute("global_seed_generator")
-    global_seed = global_state.get_global_attribute("global_random_seed")
-    if gen is None and global_seed is None:
+    if gen is None:
         gen = SeedGenerator()
         global_state.set_global_attribute("global_seed_generator", gen)
-    elif gen is None and global_seed is not None:
-        gen = SeedGenerator(global_seed)
-        global_state.set_global_attribute("global_seed_generator", gen)
     return gen
 
 

keras/src/random/seed_generator_test.py

@@ -84,11 +84,8 @@ def test_jax_tracing_with_global_seed_generator(self):
         def traced_function():
             return seed_generator.global_seed_generator().next()
 
-        with self.assertRaisesRegex(
-            ValueError,
-            "When tracing a JAX function, you should only use seeded random",
-        ):
-            traced_function()
+        result = traced_function()
+        self.assertIsNotNone(result)
 
     def test_seed_generator_serialization(self):
         random_generator = seed_generator.SeedGenerator(seed=42)

keras/src/utils/jax_utils.py

@@ -1,3 +1,5 @@
+from jax import random
+
 from keras.src import backend
 
 
@@ -9,3 +11,24 @@ def is_in_jax_tracing_scope(x=None):
             if c.__name__ == "Tracer" and c.__module__.startswith("jax"):
                 return True
     return False
+
+
+def get_jax_random_seed(seed=None):
+    if is_in_jax_tracing_scope():
+        # Constant dummy seed for Tracing
+        seed = 0
+    else:
+        # Gathering seed from a seed generator
+        seed = backend.random.draw_seed(None)[0]
+    return seed
+
+
+# Create a lightweight class that only provides shape/dtype info
+class JAXTracingSeedGenerator:
+    def __init__(self):
+        self._shape = (2,)
+        self._dtype = "uint32"
+
+    def next(self, ordered=False):
+        # Return a dummy key for tracing
+        return random.key(0)

keras/src/utils/rng_utils.py

@@ -5,6 +5,7 @@
 from keras.src import backend
 from keras.src.api_export import keras_export
 from keras.src.backend.common import global_state
+from keras.src.random import seed_generator
 from keras.src.utils.module_utils import tensorflow as tf
 
 GLOBAL_RANDOM_SEED = "global_random_seed"
@@ -60,6 +61,10 @@ def set_random_seed(seed):
         import torch
 
         torch.manual_seed(seed)
+    if backend.backend() == "jax":
+        # We create a global seed generator using the global random seed
+        gen = seed_generator.SeedGenerator(seed)
+        global_state.set_global_attribute("global_seed_generator", gen)
 
 
 def get_random_seed():