Cast in/out shardings to tuple before passing to Exported ctor.

PiperOrigin-RevId: 586951567
jax-ml · Dec 1, 2023 · 54fee48 · 54fee48
1 parent e60aa3b
commit 54fee48
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/experimental/export/export.py b/jax/experimental/export/export.py
@@ -527,8 +527,8 @@ def do_export(*args_specs, **kwargs_specs) -> Exported:
         out_tree=lowered.out_tree,
         in_avals=tuple(args_avals_flat),
         out_avals=tuple(out_avals_flat),
-        in_shardings=lowering.compile_args["in_shardings"],
-        out_shardings=lowering.compile_args["out_shardings"],
+        in_shardings=tuple(lowering.compile_args["in_shardings"]),
+        out_shardings=tuple(lowering.compile_args["out_shardings"]),
         lowering_platforms=actual_lowering_platforms,
         ordered_effects=ordered_effects,
         unordered_effects=unordered_effects,