minor fixes in tests

keras-team · Oct 16, 2024 · a8ee705 · a8ee705
1 parent 2e1d9bc
commit a8ee705
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 12 deletions.
diff --git a/keras_hub/src/models/gemma/gemma_backbone_test.py b/keras_hub/src/models/gemma/gemma_backbone_test.py
@@ -74,11 +74,10 @@ def test_architecture_characteristics(self):
 
     def test_distribution(self):
         if keras.backend.backend() != "jax":
-            return
+            self.skipTest("`ModelParallel` testing requires the Jax backend.")
         devices = keras.distribution.list_devices("CPU")
         if len(devices) == 1:
-            # Need more than 1 device for distribution testing.
-            return
+            self.skipTest("`ModelParallel` testing requires multiple devices.")
         device_mesh = keras.distribution.DeviceMesh(
             shape=(1, len(devices)),
             axis_names=("batch", "model"),
@@ -129,7 +128,6 @@ def test_distribution_with_lora(self):
             self.skipTest("`ModelParallel` testing requires the Jax backend.")
         devices = keras.distribution.list_devices("CPU")
         if len(devices) == 1:
-            # Need more than 1 device for distribution testing.
             self.skipTest("`ModelParallel` testing requires multiple devices.")
         device_mesh = keras.distribution.DeviceMesh(
             shape=(1, len(devices)),

diff --git a/keras_hub/src/models/llama/llama_backbone.py b/keras_hub/src/models/llama/llama_backbone.py
@@ -192,13 +192,20 @@ def get_layout_map(
         ```
         # Feel free to change the mesh shape to balance data and model parallelism
         mesh = keras.distribution.DeviceMesh(
-            shape=(1, 8), axis_names=('batch', 'model'),
-            devices=keras.distribution.list_devices())
+            shape=(1, 8),
+            axis_names=('batch', 'model'),
+            devices=keras.distribution.list_devices(),
+        )
         layout_map = LlamaBackbone.get_layout_map(
-            mesh, model_parallel_dim_name="model")
+            mesh,
+            model_parallel_dim_name="model",
+        )
 
         distribution = keras.distribution.ModelParallel(
-            layout_map=layout_map, batch_dim_name='batch')
+            layout_map=layout_map,
+            batch_dim_name='batch',
+        )
+
         with distribution.scope():
            llama_model = keras_hub.models.LlamaCausalLM.from_preset()
         ```

diff --git a/keras_hub/src/models/llama/llama_backbone_test.py b/keras_hub/src/models/llama/llama_backbone_test.py
@@ -1,5 +1,5 @@
-import pytest
 import keras
+import pytest
 from keras import ops
 
 from keras_hub.src.models.llama.llama_backbone import LlamaBackbone
@@ -70,11 +70,10 @@ def test_all_presets(self):
 
     def test_distribution(self):
         if keras.backend.backend() != "jax":
-            return
+            self.skipTest("`ModelParallel` testing requires the Jax backend.")
         devices = keras.distribution.list_devices("CPU")
         if len(devices) == 1:
-            # Need more than 1 device for distribution testing.
-            return
+            self.skipTest("`ModelParallel` testing requires multiple devices.")
         device_mesh = keras.distribution.DeviceMesh(
             shape=(1, len(devices)),
             axis_names=("batch", "model"),