Internal Change

KfacJaxDev · KfacJaxDev · commit 8c48c95e53f3 · 2025-09-16T04:51:37.000-07:00
PiperOrigin-RevId: 805859317
diff --git a/kfac_jax/_src/utils/staging.py b/kfac_jax/_src/utils/staging.py
@@ -142,7 +142,7 @@ def pmean_if_pmap_wrapper(
       func: Callable[..., TArrayTree],
   ) -> Callable[..., TArrayTree]:
     """Wraps a function to perform a pmean if `multi_device`."""
-    if self.multi_device:
+    if self.multi_device and not self.debug:
       return lambda *args, **kwargs: lax.pmean(
           func(*args, **kwargs), self.pmap_axis_name
       )

Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,7 @@ def pmean_if_pmap_wrapper(`
`142`	`142`	`func: Callable[..., TArrayTree],`
`143`	`143`	`) -> Callable[..., TArrayTree]:`
`144`	`144`	"""Wraps a function to perform a pmean if `multi_device`."""
`145`		`- if self.multi_device:`
	`145`	`+ if self.multi_device and not self.debug:`
`146`	`146`	`return lambda args, *kwargs: lax.pmean(`
`147`	`147`	`func(args, *kwargs), self.pmap_axis_name`
`148`	`148`	`)`