Reusing computation results across multiple gradients in JAX

[itk22]

Hello,
I'm working on a project where I need to compute two separate values (an objective and a constraint) and their gradients, both of which depend on the same expensive operation. Currently, when I compute the gradients separately using jax.grad, the expensive operation is executed twice (see the example below). Given the high computational cost of that operation, I would like to avoid repeating it.

import jax
from jax import grad, jit
import jax.numpy as jnp

# Dummy expensive operation
def expensive_function(x):
    return jnp.square(x)

# Objective
def objective_fn(x):
    y = expensive_function(x)
    return jnp.sum(y)

# Constraint
def constraint_fn(x):
    y = expensive_function(x)
    return jnp.sum(y * 2)  

# Inputs
x = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=jnp.float32)

# Gradients
objective_grad = grad(objective_fn)(x)
constraint_grad = grad(constraint_fn)(x)

Is there a way to compute the gradients of both values together, while reusing the result of the expensive operation? I understand that JAX avoids side effects and mutable state by design, which makes caching results across multiple function calls tricky. Is there any workaround or recommended way to handle this type of scenario?

Or should I perhaps rely on JIT for common subexpression elimination as outlined in Discussion #16356?

[davisyoshida]

If you want to put stuff together manually in this case you could use jax.vjp, but you'd need to refactor objective_fn and constraint_fn to take y as an argument. Here's what it might look like:

import jax    
from jax import grad, jit    
import jax.numpy as jnp    
    
# Dummy expensive operation    
def expensive_function(x):    
    return jnp.square(x)    
    
# Objective    
def objective_fn(y):    
    return jnp.sum(y)    
    
# Constraint    
def constraint_fn(y):    
    return jnp.sum(y * 2)··    
    
# Inputs    
x = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=jnp.float32)    
    
# Gradients    
y, vjp_fn = jax.vjp(expensive_function, x)    
objective_ygrad = grad(objective_fn)(y)    
constraint_ygrad = grad(constraint_fn)(y)    
    
objective_grad = vjp_fn(objective_ygrad)[0]    
constraint_grad = vjp_fn(constraint_ygrad)[0]

I'm not 100% sure this is correct since I'm not a JAX expert, but I believe it's basically doing what you want. The Autodiff Cookbook is a great resource which explains the various autodiff options in JAX.

[itk22]

Thank you for your answer @davisyoshida. This looks promising and I will make sure to check out the autodiff cookbook in more detail!

[davisyoshida]

Happy to help! BTW I do think it's likely that CSE will do the trick if you jit everything together, but if you need to be jitting things separately for whatever reason this may help.