Optimizations w.r.t. redundant gradient computations

[mchagneux]

Hi all,

I'm working on code which involves taking derivatives w.r.t some parameter $\theta$ of composite functions of the form

$$L(\theta) = f(\psi(\theta)) \times g(\psi(\theta))$$

For some reasons which would be out of topic to describe here I cannot rely on autodifferentiation alone. Writing $F(\theta) = f(\psi(\theta))$ and $G(\theta) = g(\psi(\theta))$, I need to explicitely compute

$$\nabla L = F \times \nabla G + \nabla F \times G$$

However I can use autodifferentiation for $F$ and $G$ separately, ie. I can define jax.grad(g) and jax.grad(f) and write something like this

@jit
def grad_L(theta):

 return F(theta) * jax.grad(G)(theta) + jax.grad(F)(theta) * G(theta)

Since the gradients of $F$ and $G$ both involve the Jacobian of $\psi$, will the compiler optimize this code such that this Jacobian is only computed once ? Or should I rather compute this Jacobian myself via jax.jacrev or jax.jacfwd and compute the gradients of $F$ and $G$ via the chain rule written by hand ?

Thanks in advance!

[jakevdp]

Without more details it's hard to say how the compiler will optimize any particular sequence of operations, but in general it should. The most definitive way to check is to print the compiled HLO to see. Here's a simple example:

import jax
import jax.numpy as jnp

def f(x):
  y = jnp.sin(x)
  z = jnp.sin(x)
  return y + z

print(jax.jit(f).lower(1.0).compile().as_text())

HloModule jit_f, entry_computation_layout={(f32[])->f32[]}

%fused_computation (param_0.1: f32[]) -> f32[] {
  %param_0.1 = f32[] parameter(0)
  %sine.0 = f32[] sine(f32[] %param_0.1), metadata={op_name="jit(f)/jit(main)/sin" source_file="<ipython-input-34-90031d99c954>" source_line=5}
  ROOT %add.0 = f32[] add(f32[] %sine.0, f32[] %sine.0), metadata={op_name="jit(f)/jit(main)/add" source_file="<ipython-input-34-90031d99c954>" source_line=7}
}

ENTRY %main.5 (Arg_0.1: f32[]) -> f32[] {
  %Arg_0.1 = f32[] parameter(0)
  ROOT %fusion = f32[] fusion(f32[] %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(f)/jit(main)/add" source_file="<ipython-input-34-90031d99c954>" source_line=7}
}

You can see that in this case, the sine is only computed once (%sine.0 = f32[] sine(f32[] %param_0.1)), and passed to both sides of the add operation (%add.0 = f32[] add(f32[] %sine.0, f32[] %sine.0)).

[mchagneux]

Ok, thanks for the tip I wasn't aware of this feature!