Only enable JAX on linux_x86_64. (#101)

Author: cantonios

keras_rs/src/layers/embedding/distributed_embedding.py

@@ -1,8 +1,16 @@
+import platform
+import sys
+
 import keras
 
 from keras_rs.src.api_export import keras_rs_export
 
-if keras.backend.backend() == "jax":
+# JAX TPU embedding is only available on linux_x86_64.
+if (
+    keras.backend.backend() == "jax"
+    and sys.platform == "linux"
+    and platform.machine().lower() == "x86_64"
+):
     from keras_rs.src.layers.embedding.jax.distributed_embedding import (
         DistributedEmbedding as BackendDistributedEmbedding,
     )

requirements-jax-cuda.txt

@@ -11,6 +11,6 @@ torch>=2.1.0
 jax[cuda12_pip]==0.6.0
 
 # Support for large embeddings.
-jax-tpu-embedding
+jax-tpu-embedding;sys_platform == 'linux' and platform_machine == 'x86_64'
 
 -r requirements-common.txt

requirements.txt

@@ -8,7 +8,7 @@ torch>=2.1.0
 
 # Jax.
 jax[cpu]
-jax-tpu-embedding
+jax-tpu-embedding;sys_platform == 'linux' and platform_machine == 'x86_64'
 
 # pre-commit checks (formatting, linting, etc.)
 pre-commit