experimental multi-host device_put

Author: Lumosis

tpu_inference/utils.py

@@ -243,7 +243,7 @@ def device_array(mesh: Mesh, *args, sharding=None, **kwargs) -> jax.Array:
     """
     if sharding is None:
         sharding = NamedSharding(mesh, PartitionSpec(None))
-    return jax.make_array_from_process_local_data(sharding=sharding, *args, **kwargs)
+    return jax.make_array_from_process_local_data(sharding, *args, **kwargs)
 
 
 def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]: