Fix ModelParallel sharded variable loading test

- Change Embedding input_dim from 100 to 96 to be divisible by 8 devices
- Add skip condition for tests requiring < 2 devices (sharding not meaningful)
- Prevents ValueError when sharding dimensions not divisible by device count
- Test now properly skips on single-device systems and passes on multi-device CI

Author: amitsrivastava78

keras/src/distribution/distribution_lib_test.py

@@ -381,6 +381,11 @@ def test_model_parallel_sharded_variable_loading(self):
 
         # Ensure we have JAX devices
         jax_devices = jax.devices()
+        if len(jax_devices) < 2:
+            pytest.skip(
+                "Test requires at least 2 devices for meaningful sharding"
+            )
+
         logging.debug(f"JAX devices available: {len(jax_devices)}")
         for i, device in enumerate(jax_devices):
             logging.debug(f"  Device {i}: {device}")
@@ -432,7 +437,7 @@ def test_model_parallel_sharded_variable_loading(self):
                     ),
                     # Embedding layer (modified in commit)
                     layers.Embedding(
-                        input_dim=100, output_dim=32, name="embedding"
+                        input_dim=96, output_dim=32, name="embedding"
                     ),
                     layers.Flatten(),
                     # Convolutional layer (modified in commit)
@@ -517,7 +522,7 @@ def test_model_parallel_sharded_variable_loading(self):
                             "ab,bc->ac", output_shape=32, name="einsum_dense"
                         ),
                         layers.Embedding(
-                            input_dim=100, output_dim=32, name="embedding"
+                            input_dim=96, output_dim=32, name="embedding"
                         ),
                         layers.Flatten(),
                         layers.Reshape(