Fix handling of literals for loop

  - d55f103e9518b6394ae4879c408dd055e44e3e04 by Anxhelo Xhebraj <axhebraj@nvidia.com>

Signed-off-by: Anxhelo Xhebraj <axhebraj@nvidia.com>
GitOrigin-RevId: d55f103e9518b6394ae4879c408dd055e44e3e04

Author: Angelogeb

src/jaxpp/training.py

@@ -22,9 +22,12 @@
 import jax._src.core as jcore
 import jax.api_util as jau
 import jax.extend.linear_util as lu
-import jax.numpy as jnp
 import numpy as np
 from jax._src import dtypes
+
+# TODO: maybe use custom definition of `tree_broadcast` and
+#  improve error message if it cannot be broadcasted
+from jax._src.custom_transpose import tree_broadcast
 from jax.interpreters import ad
 from jax.interpreters import partial_eval as pe
 
@@ -128,6 +131,7 @@ class Concat:
     axis: int = 0
 
     def state(self, n: int, a: jax.ShapeDtypeStruct) -> jax.Array:
+        assert n > 0
         shape = a.shape[: self.axis] + (n,) + a.shape[self.axis :]
         return jax.numpy.zeros(shape, dtype=a.dtype)
 
@@ -218,14 +222,20 @@ def treduce(fun, xs, operation=(Concat(), Add())):
 
     @functools.wraps(fun)
     def wrap(i):
-        e = jax.tree.map(lambda x: jnp.take(x, i, axis=axis), xs)
+        e = jax.tree.map(lambda x: jax.numpy.take(x, i, axis=axis), xs)
         return fun(e)
 
     return treduce_i(
         wrap, first_batch_shape[axis], schedule=schedule, operation=operation
     )
 
 
+def copy_if_scalar(x: jax.Array) -> jax.Array:
+    if x.ndim == 0:
+        return jax.numpy.array(x, copy=True)
+    return x
+
+
 def treduce_i(
     fun: Callable[[int], Y], length: int, schedule: BaseSchedule, operation=default_op
 ) -> Y:
@@ -264,25 +274,19 @@ def treduce_i(fun, length, operation):
       structure as ``Y``.
     """
     with log_elapsed_time("jaxpr/first_loop_tracing"), yield_scope():
-        body_args = jcore.ShapedArray((), dtype=jnp.int32)
-        vmapped_jaxpr, loop_out_shapes = jax.make_jaxpr(fun, return_shape=True)(
-            body_args
-        )
-
-    # TODO: maybe use custom definition of `tree_broadcast` and
-    #  improve error message if it cannot be broadcasted
-    from jax._src.custom_transpose import tree_broadcast
+        body_args = jcore.ShapedArray((), dtype=jax.numpy.int32)
+        body_jaxpr, loop_out_shapes = jax.make_jaxpr(fun, return_shape=True)(body_args)
 
     operation = tree_broadcast(jax.tree_util.tree_structure(loop_out_shapes), operation)
 
     def state(op: Op, a):
-        return op.state(length, a)
+        return copy_if_scalar(op.state(length, a))
 
     loop_state = jax.tree_util.tree_map(state, operation, loop_out_shapes)
 
     def _fun(mubatch_idx, loop_state):
         def update(op: Op, state, update):
-            return op.update(state, update, mubatch_idx)
+            return copy_if_scalar(op.update(state, update, mubatch_idx))
 
         return (
             mubatch_idx + 1,
@@ -293,8 +297,8 @@ def update(op: Op, state, update):
                 jax.tree.unflatten(
                     jax.tree.structure(loop_out_shapes),
                     jcore.eval_jaxpr(
-                        vmapped_jaxpr.jaxpr,
-                        vmapped_jaxpr.consts,
+                        body_jaxpr.jaxpr,
+                        body_jaxpr.consts,
                         mubatch_idx,
                         propagate_source_info=False,
                     ),
@@ -303,10 +307,10 @@ def update(op: Op, state, update):
         )
 
     debug_info = jau.debug_info(treduce_i.__name__, fun, (body_args,), {})
-    wrapped_vmapped_fun = lu.wrap_init(_fun, debug_info=debug_info)
+    wrapped_body_fun = lu.wrap_init(_fun, debug_info=debug_info)
     with log_elapsed_time("jaxpr/second_loop_tracing"), yield_scope():
         loop_output = pscan_wrapped(
-            wrapped_vmapped_fun, loop_state, length=length, schedule=schedule
+            wrapped_body_fun, loop_state, length=length, schedule=schedule
         )
 
     return loop_output

tests/test_transformations.py

@@ -11,11 +11,14 @@
 from jax.scipy.special import logsumexp
 
 from jaxpp import env_vars
-from jaxpp.api import BaseSchedule, pipeline_enter_stage, treduce
+from jaxpp.api import Add, BaseSchedule, Concat, pipeline_enter_stage, treduce
 from jaxpp.core import (
     cluster_jaxpr,
+    common_passes,
+    fixup_multidefs,
     infer_shardings2,
     maybe_unroll_loop,
+    outvar_normalization,
     strip_inspect_sharding_eqns,
     wrap_into_tasks,
 )
@@ -213,6 +216,42 @@ def test_transformations_dont_fail(
     )
 
 
+def test_literal_via_constant_folding():
+    """Test if constant folding or other optimizations might introduce literals."""
+    num_stages = 1
+    n_mubatches = 2
+    schedule = Std1F1B(num_stages=1)
+    (params, X, Y), stage_mesh, replicated_sharding = get_context(
+        num_stages, n_mubatches
+    )
+
+    def grads_with_constant_output(params, data):
+        loss, grad = grads(params, data)
+        identity = jnp.array(1.0)
+        return (loss, grad, identity)
+
+    custom_op = (Concat(), Add, Add)
+
+    total_grads_fn = jax.jit(
+        lambda params, X, Y: treduce(
+            partial(grads_with_constant_output, params),
+            (X, Y),
+            schedule=schedule,
+            operation=custom_op,
+        )
+    )
+
+    cjaxpr = total_grads_fn.trace(params, X, Y).jaxpr
+    with env_vars.jaxpp_conservative_loop_clustering.set(False):
+        scheduled_jaxpr = get_scheduled_jaxpr(
+            cjaxpr,
+            1,
+            stage_mesh,
+            replicated_sharding,
+            skip_propagation=True,
+        )
+
+
 def test_skip_propagation_false():
     test_equivalence_scheduled(
         *(3, 3, 5, Eager1F1B(num_stages=3)), skip_propagation=False