Check for jax_tpu_embedding on JAX backend. (#114)

This is to allow users to potentially run Keras RS _without_ the
dependency.

If a user doesn't have `jax-tpu-embedding` installed, but are on
`linux_x86_64` and has a sparsecore-capable TPU available, and
if they try to use `auto` or `sparsecore` placement with distributed
embedding, will raise an error informing them to install the
dependency.

Author: cantonios

keras_rs/src/layers/embedding/base_distributed_embedding.py

@@ -1,4 +1,5 @@
 import collections
+import importlib.util
 import typing
 from typing import Any, Sequence
 
@@ -304,7 +305,7 @@ def step(data):
     To use `DistributedEmbedding` on TPUs with JAX, one must create and set a
     Keras `Distribution`.
     ```
-    distribution = keras.distribution.DataParallel(devices=jax.device("tpu))
+    distribution = keras.distribution.DataParallel(devices=jax.device("tpu"))
     keras.distribution.set_distribution(distribution)
     ```
 
@@ -901,6 +902,24 @@ def _default_device_get_embedding_tables(self) -> dict[str, types.Tensor]:
         return tables
 
     def _has_sparsecore(self) -> bool:
+        # Explicitly check for SparseCore availability.
+        # We need this check here rather than in jax/distributed_embedding.py
+        # so that we can warn the user about missing dependencies.
+        if keras.backend.backend() == "jax":
+            # Check if SparseCores are available.
+            try:
+                import jax
+
+                tpu_devices = jax.devices("tpu")
+            except RuntimeError:
+                # No TPUs available.
+                return False
+
+            if len(tpu_devices) > 0:
+                device_kind = tpu_devices[0].device_kind
+                if device_kind in ["TPU v5", "TPU v6 lite"]:
+                    return True
+
         return False
 
     def _sparsecore_init(
@@ -909,6 +928,17 @@ def _sparsecore_init(
         table_stacking: str | Sequence[Sequence[str]],
     ) -> None:
         del feature_configs, table_stacking
+
+        if keras.backend.backend() == "jax":
+            jax_tpu_embedding_spec = importlib.util.find_spec(
+                "jax_tpu_embedding"
+            )
+            if jax_tpu_embedding_spec is None:
+                raise ImportError(
+                    "Please install jax-tpu-embedding to use "
+                    "DistributedEmbedding on sparsecore devices."
+                )
+
         raise self._unsupported_placement_error("sparsecore")
 
     def _sparsecore_build(self, input_shapes: dict[str, types.Shape]) -> None:

keras_rs/src/layers/embedding/distributed_embedding.py

@@ -1,15 +1,19 @@
+import importlib.util
 import platform
 import sys
 
 import keras
 
 from keras_rs.src.api_export import keras_rs_export
 
-# JAX TPU embedding is only available on linux_x86_64.
+# JAX distributed embedding is only available on linux_x86_64, and only if
+# jax-tpu-embedding is installed.
+jax_tpu_embedding_spec = importlib.util.find_spec("jax_tpu_embedding")
 if (
     keras.backend.backend() == "jax"
     and sys.platform == "linux"
     and platform.machine().lower() == "x86_64"
+    and jax_tpu_embedding_spec is not None
 ):
     from keras_rs.src.layers.embedding.jax.distributed_embedding import (
         DistributedEmbedding as BackendDistributedEmbedding,

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

@@ -347,12 +347,6 @@ def _add_table_variable(
 
         return table_variable, slot_variables
 
-    def _has_sparsecore(self) -> bool:
-        device_kind = jax.devices()[0].device_kind
-        if device_kind in ["TPU v5", "TPU v6 lite"]:
-            return True
-        return False
-
     @keras_utils.no_automatic_dependency_tracking
     def _sparsecore_init(
         self,