Running Partitioner.compile within Mesh context-manager

google-research · Nov 22, 2023 · d08a684 · d08a684
1 parent 4b8a1ad
commit d08a684
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/t5x/partitioning.py b/t5x/partitioning.py
@@ -783,6 +783,13 @@ def lower(self, *args, **kwargs):
                   self._logical_axis_rules):
       return self._pjitted_fn.lower(*args, **kwargs)
 
+  def lower_and_compile(self, *args, **kwargs):
+    with Mesh(self._mesh.devices,
+              self._mesh.axis_names), flax_partitioning.axis_rules(
+                  self._logical_axis_rules):
+      return self._pjitted_fn.lower(*args, **kwargs).compile()
+
+
 
 class BasePjitPartitioner(BasePartitioner):
   """Partitioner that uses T5X version of jax.pjit."""
@@ -816,7 +823,7 @@ def partition(
 
   def compile(self, partitioned_fn: PjittedFnWithContext,
               *args) -> CompiledPartitionedCallable:
-    return partitioned_fn.lower(*args).compile()
+    return partitioned_fn.lower_and_compile(*args)
 
 
 class PjitPartitioner(BasePjitPartitioner):