Ignore shard_map attr error in mypy. (#97)

Author: cantonios

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

@@ -30,7 +30,7 @@
 
 ArrayLike = Union[np.ndarray[Any, Any], jax.Array]
 FeatureConfig = config.FeatureConfig
-shard_map = jax.experimental.shard_map.shard_map
+shard_map = jax.experimental.shard_map.shard_map  # type: ignore[attr-defined]
 
 
 def _get_partition_spec(

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

@@ -16,7 +16,7 @@
 from keras_rs.src.layers.embedding.jax import embedding_utils
 
 ShardedCooMatrix = embedding_utils.ShardedCooMatrix
-shard_map = jax.experimental.shard_map.shard_map
+shard_map = jax.experimental.shard_map.shard_map  # type: ignore[attr-defined]
 
 T = TypeVar("T")
 Nested = Union[T, Sequence[T], Mapping[str, T]]