fixed bug with gather functions

young-geng · Jan 7, 2024 · 2f3df6d · 2f3df6d
1 parent b444ed0
commit 2f3df6d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scalax/sharding.py b/scalax/sharding.py
@@ -188,7 +188,7 @@ def wrapped(*args):
             else:
                 output = jax.eval_shape(lambda: fun(*args))
                 matched_out_shardings = self._match_sharding_rule(out_shardings, output)
-
+            
             jitted_fn = jax.jit(
                 sharding_constrained_fun,
                 in_shardings=matched_in_shardings,
@@ -233,7 +233,7 @@ def make_gather_fn(partition_spec):
             jax_gather_fn = jax.jit(
                 lambda x: x,
                 in_shardings=partition_spec,
-                out_shardings=None
+                out_shardings=NamedSharding(self.mesh, PartitionSpec()),
             )
             def gather_fn(tensor):
                 return jax.device_get(jax_gather_fn(tensor))