made distributed structure proper

Author: amitsrivastava78

keras/src/backend/__init__.py

@@ -36,11 +36,9 @@
 # Import backend functions.
 if backend() == "tensorflow":
     from keras.src.backend.tensorflow import *  # noqa: F403
-    from keras.src.backend.tensorflow import distribution_lib
     from keras.src.backend.tensorflow.core import Variable as BackendVariable
 elif backend() == "jax":
     from keras.src.backend.jax import *  # noqa: F403
-    from keras.src.backend.jax import distribution_lib
     from keras.src.backend.jax.core import Variable as BackendVariable
 elif backend() == "torch":
     from keras.src.backend.torch import *  # noqa: F403

keras/src/backend/jax/__init__.py

@@ -1,5 +1,6 @@
 from keras.src.backend.config import is_nnx_enabled
 from keras.src.backend.jax import core
+from keras.src.backend.jax import distribution_lib
 from keras.src.backend.jax import image
 from keras.src.backend.jax import linalg
 from keras.src.backend.jax import math

keras/src/backend/jax/distribution_lib.py

@@ -193,6 +193,11 @@ def num_processes():
     return jax.process_count()
 
 
+def process_id():
+    """Return the current process ID for the distribution setting."""
+    return jax.process_index()
+
+
 def _to_backend_device(device_name):
     if isinstance(device_name, jax.Device):
         return device_name

keras/src/backend/numpy/__init__.py

@@ -24,6 +24,3 @@
 from keras.src.backend.numpy.rnn import gru
 from keras.src.backend.numpy.rnn import lstm
 from keras.src.backend.numpy.rnn import rnn
-
-# Numpy backend does not support distribution
-distribution_lib = None

keras/src/backend/tensorflow/__init__.py

@@ -1,4 +1,5 @@
 from keras.src.backend.tensorflow import core
+from keras.src.backend.tensorflow import distribution_lib
 from keras.src.backend.tensorflow import image
 from keras.src.backend.tensorflow import linalg
 from keras.src.backend.tensorflow import math

keras/src/wrappers/sklearn_test.py

@@ -120,6 +120,9 @@ def use_floatx(x):
             "not an issue in sklearn>=1.6"
         ),
         "check_pipeline_consistency": "Neural networks are non-deterministic",
+        "check_transformer_data_not_an_array": "Neural networks are "
+        "non-deterministic",
+        "check_transformer_general": "Neural networks are non-deterministic",
     },
 }