Does JIT avoid recomputing gradients?

[mchagneux]

Hi all,

I was about the behaviour of JAX relative to the chain rule when I want to perform certain operations on the gradients. Say I have the following code:

@jax.jit
def foo(theta):
 
 def f(theta):
  return <some-thing-that-depends-on-theta> 

 def g(theta):
  return f(theta) + <some-other-thing-that-depends-on-theta>
 
 def h(theta):
  return f(theta) + <yet-another-thing-that-depends-on-theta>

return tree_map(lambda x,y: x*y, jax.grad(g)(theta), jax.grad(h)(theta))

Can I be certain that JAX won't compute the gradient of f twice ?

[davisyoshida]

The JIT does common subexpression elimination. Here's an example:

import jax      
import jax.numpy as jnp      
      
def foo(x):      
    def f(x):      
        return jnp.sin(x)      
      
    def g(x):      
        return f(x) + x      
      
    def h(x):      
        return f(x) + 1/x      
      
    return jax.grad(g)(x), jax.grad(h)(x)      
      
jitted = jax.jit(foo)      
compiled = jitted.lower(1.).compile()      
print(compiled.as_text())

This outputs:

HloModule jit_foo, is_scheduled=true, entry_computation_layout={(f32[])->(f32[], f32[])}, allow_spmd_sharding_propagation_to_output={true,true}

%fused_computation.1 (param_0.5: f32[]) -> (f32[], f32[]) {
  %param_0.5 = f32[] parameter(0)
  %cosine.4 = f32[] cosine(f32[] %param_0.5), metadata={op_name="jit(foo)/jit(main)/cos" source_file="/home/davis/src/scratch/jax/inspect_jaxpr.py" source_line=6}
  %constant_4 = f32[] constant(1)
  %add.3 = f32[] add(f32[] %cosine.4, f32[] %constant_4), metadata={op_name="jit(foo)/jit(main)/add_any" source_file="/home/davis/src/scratch/jax/inspect_jaxpr.py" source_line=6}
  %multiply.1.clone.1 = f32[] multiply(f32[] %param_0.5, f32[] %param_0.5), metadata={op_name="jit(foo)/jit(main)/mul" source_file="/home/davis/src/scratch/jax/inspect_jaxpr.py" source_line=12}
  %divide.1.clone.1 = f32[] divide(f32[] %constant_4, f32[] %multiply.1.clone.1), metadata={op_name="jit(foo)/jit(main)/div" source_file="/home/davis/src/scratch/jax/inspect_jaxpr.py" source_line=12}
  %negate.1.clone.1 = f32[] negate(f32[] %divide.1.clone.1), metadata={op_name="jit(foo)/jit(main)/neg" source_file="/home/davis/src/scratch/jax/inspect_jaxpr.py" source_line=12}
  %add.2.clone.1 = f32[] add(f32[] %negate.1.clone.1, f32[] %cosine.4), metadata={op_name="jit(foo)/jit(main)/add_any" source_file="/home/davis/src/scratch/jax/inspect_jaxpr.py" source_line=6}
  ROOT %tuple = (f32[], f32[]) tuple(f32[] %add.3, f32[] %add.2.clone.1)
}

ENTRY %main.11 (Arg_0.1: f32[]) -> (f32[], f32[]) {
  %Arg_0.1 = f32[] parameter(0), sharding={replicated}
  ROOT %fusion.1 = (f32[], f32[]) fusion(f32[] %Arg_0.1), kind=kLoop, calls=%fused_computation.1, frontend_attributes={fingerprint_before_lhs="133e643e0b721ac98faebb3592b66bf2"}, metadata={op_name="jit(foo)/jit(main)/add_any" source_file="/home/davis/src/scratch/jax/inspect_jaxpr.py" source_line=6}
}

As you can see there's only one call to cosine(), so at least in this case it's successfully re-using that computation. I'd expect it to do so in general but perhaps the devs will know of some cases where the compiler can't figure it out.

[mchagneux]

Ok thanks! Indeed mostly wondering about how systematic it is. In many aspects of my work I'm only interested in gradients but it's really more convenient to write multiple calls to intermediate functions because I might need both the gradient of f and the gradient of something that depends on the output of f. I've been doing this kind of thing everywhere in my code and strongly hoping that JIT does the right work here because at some point it becomes impossible to check the IR :)

[mchagneux]

Does anyone from the JAX development team has more insight into this ? When I profile my code with stack traces I see that gradients of functions seem to recompute the gradients of the inner functions they have in common. Is this just an artifact of the way stack traces are displayed ? Lowered expressions are impossible to read for my code.