fix test errors

Author: wenyi-guo

keras_rs/src/layers/embedding/jax/distributed_embedding_test.py

@@ -28,6 +28,7 @@
 
 keras.config.disable_traceback_filtering()
 
+from keras_rs.src import testing
 
 def _create_sparsecore_layout(
     sharding_axis: str = "sparsecore",
@@ -308,7 +309,7 @@ def my_initializer(shape: tuple[int, int], dtype: Any):
     keras.backend.backend() != "jax",
     reason="Backend specific test",
 )
-class DistributedEmbeddingLayerTest(parameterized.TestCase):
+class DistributedEmbeddingLayerTest(testing.TestCase, parameterized.TestCase):
     @parameterized.product(
         ragged=[True, False],
         combiner=["sum", "mean", "sqrtn"],
@@ -326,6 +327,7 @@ def test_call(
         table_stacking: str | list[str] | list[list[str]],
         jit: bool,
     ):
+        self.on_tpu = "TPU_NAME" in os.environ
         table_configs = keras_test_utils.create_random_table_configs(
             combiner=combiner, seed=10
         )
@@ -374,7 +376,7 @@ def test_call(
         )
 
         keras.tree.map_structure(
-            lambda a, b: np.testing.assert_allclose(a, b, atol=1e-5),
+            lambda a, b: self.assertAllClose(a, b, atol=1e-3, is_tpu=self.on_tpu),
             outputs,
             expected_outputs,
         )

keras_rs/src/utils/tpu_test_utils.py

@@ -1,7 +1,7 @@
 import contextlib
 import os
 from types import ModuleType
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Any, Callable, ContextManager, Optional, Tuple, Union
 
 import keras
 import tensorflow as tf
@@ -15,24 +15,26 @@
 
 
 class DummyStrategy:
-    def scope(self):
+    def scope(self) -> ContextManager[None]:
         return contextlib.nullcontext()
 
     @property
-    def num_replicas_in_sync(self):
+    def num_replicas_in_sync(self) -> int:
         return 1
 
-    def run(self, fn, args):
+    def run(self, fn: Callable[..., Any], args: Tuple[Any, ...]) -> Any:
         return fn(*args)
 
-    def experimental_distribute_dataset(self, dataset, options=None):
+    def experimental_distribute_dataset(
+        self, dataset: Any, options: Optional[Any] = None
+    ) -> Any:
         del options
         return dataset
 
 
 class JaxDummyStrategy(DummyStrategy):
     @property
-    def num_replicas_in_sync(self):
+    def num_replicas_in_sync(self) -> int:
         if jax is None:
             return 0
         return jax.device_count("tpu")
@@ -87,7 +89,7 @@ def run_with_strategy(
         sample_weight_value = kwargs.get("sample_weight", None)
         all_inputs = args + (sample_weight_value,)
 
-        @tf.function(jit_compile=jit_compile)
+        @tf.function(jit_compile=jit_compile)  # type: ignore[misc]
         def tf_function_wrapper(input_tuple: Tuple[Any, ...]) -> Any:
             num_original_args = len(args)
             core_args = input_tuple[:num_original_args]