Merge pull request #25048 from jax-ml:linearization-rule-signature

PiperOrigin-RevId: 699007033

Author: Google-ML-Automation

jax/_src/interpreters/ad.py

@@ -483,39 +483,44 @@ def to_primal_tangent_pair(self, val):
 
   def process_primitive(self, primitive, args, params):
     primals_in, tangents_in = unzip2(map(self.to_primal_tangent_pair, args))
+    tangent_nonzeros = [type(t) is not Zero for t in tangents_in]
     if all(type(t) is Zero for t in tangents_in):
       return primitive.bind_with_trace(self.parent_trace, primals_in, params)
     lin = primitive_linearizations.get(primitive)
     if lin is None:
       lin = partial(fallback_linearize_rule, primitive)
     with core.set_current_trace(self.parent_trace):
-      primal_out, linearized = lin(*primals_in, **params)
+      primal_out, tangent_nonzeros_out, residuals, linearized = lin(
+          tangent_nonzeros, *primals_in, **params)
     with core.set_current_trace(self.tangent_trace):
-      tangent_out = linearized(*tangents_in)
+      tangent_out = linearized(residuals, *tangents_in)
     if primitive.multiple_results:
-      return [maybe_linearize_tracer(self, x, t) for x, t in zip(primal_out, tangent_out)]
+      return [maybe_linearize_tracer(self, x, nz, t)
+              for x, nz, t in zip(primal_out, tangent_nonzeros, tangent_out)]
     else:
-      return maybe_linearize_tracer(self, primal_out, tangent_out)
+      return maybe_linearize_tracer(self, primal_out, tangent_nonzeros, tangent_out)
 
-def maybe_linearize_tracer(trace, primal, tangent):
-  if type(tangent) is Zero:
-    return primal
-  else:
+def maybe_linearize_tracer(trace, primal, is_nonzero, tangent):
+  if is_nonzero:
+    assert not type(tangent) is Zero
     return LinearizeTracer(trace, primal, tangent)
+  else:
+    assert type(tangent) is Zero
+    return primal
 
-def fallback_linearize_rule(prim, *args, **kwargs):
+def fallback_linearize_rule(prim, _, *args, **kwargs):
   def call_prim(*args_):
     return prim.bind(*args_, **kwargs)
   with config.use_direct_linearize(False):
     out_primals, out_tangents_pvals, jaxpr, consts, *_maybe_aux = linearize(
         lu.wrap_init(call_prim), *args, **kwargs)
-  def linearized(*tangents):
-    tangents_out = iter(core.eval_jaxpr(jaxpr, consts, *tangents))
+  def linearized(residuals, *tangents):
+    tangents_out = iter(core.eval_jaxpr(jaxpr, residuals, *tangents))
     full_out = [pval.get_known() if pval.is_known() else next(tangents_out)
                 for pval in out_tangents_pvals]
     assert next(tangents_out, None) is None
     return full_out
-  return out_primals, linearized
+  return out_primals, [True for _ in out_primals], consts, linearized
 
 class LinearizeTracer(Tracer):
   __slots__ = ['primal', 'tangent']
@@ -547,7 +552,7 @@ def to_concrete_value(self):
 primitive_transposes: dict[core.Primitive, Callable] = {}
 # transpose rules that internally perform reductions over the given named axes
 reducing_transposes: dict[core.Primitive, Callable] = {}
-primitive_linearizations: dict[core.Primitive, Callable]  = {}
+primitive_linearizations : dict[core.Primitive, Callable]  = {}
 
 def deflinear(primitive, transpose_rule):
   primitive_jvps[primitive] = partial(linear_jvp, primitive)

jax/_src/lax/lax.py

@@ -2400,12 +2400,16 @@ def _sin_lowering(ctx, x):
     return sine(ctx, x)
   return _nary_lower_hlo(hlo.sine, ctx, x)
 
+def _sin_p_lin(_, x):
+  cos_x = cos(x) # TODO: allow this to happen in the linearized computation (need to fix backward_pass)
+  return (sin_p.bind(x), True, cos_x, lambda cos_x_, t: mul(t, cos_x_))
+
 sin_p = standard_unop(_float | _complex, 'sin')
 ad.defjvp(sin_p, lambda g, x: mul(g, cos(x)))
+ad.primitive_linearizations[sin_p] = _sin_p_lin
 mlir.register_lowering(sin_p, _sin_lowering)
 batching.ragged_prop_rules[sin_p] = batching.ragged_mask_elementwise_rule
 
-
 def _cos_complex(x):
   # cos(x) = complex(cos(real(x)) * cosh(imag(x)), -sin(real(x)) * sinh(imag(x)))
   # see also _sin_complex

tests/api_test.py

@@ -4818,7 +4818,7 @@ def check_invariant_to_use_direct_linearize(f):
       self.assertEqual(ans1, ans2)
 
     def sin_of_sin(x):
-      return jnp.sin(jnp.sin(x))
+      return lax.sin(lax.sin(x))
 
     check_invariant_to_use_direct_linearize(lambda: jax.grad(sin_of_sin)(1.0))