Add out_sharding argument WrappedKerasInitializer. (#102)

This is for forward-compatibility.  Latest versions of JAX introduce the
`out_sharding` argument.

Author: cantonios

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

@@ -36,8 +36,13 @@ def key(self) -> Union[jax.Array, None]:
         return None
 
     def __call__(
-        self, key: Any, shape: Any, dtype: Any = jnp.float_
+        self,
+        key: Any,
+        shape: Any,
+        dtype: Any = jnp.float_,
+        out_sharding: Any = None,
     ) -> jax.Array:
+        del out_sharding
         # Force use of provided key.  The JAX backend for random initializers
         # forwards the `seed` attribute to the underlying JAX random functions.
         if key is not None and hasattr(self.initializer, "seed"):