Fix for jax import in distributed_embedding_test.py. (#162)

Passing works better in some environments.

Author: hertschuh

keras_rs/src/layers/embedding/distributed_embedding_test.py

@@ -19,8 +19,7 @@
     import jax
     import jax.experimental.sparse as jax_sparse
 except ImportError:
-    jax = None
-    jax_sparse = None
+    pass
 
 
 FEATURE1_EMBEDDING_OUTPUT_DIM = 7