custom_jvp calculations with and without jax.jit

[LawrenceMMStewart]

Hello,

I have a question about how jax will compute custom_jvps with and without jax.jit in both forward and reverse-mode differentiation.

Consider the following function forward whose computation involves calling a function expensive_computation (taking inputs x and a rng), which returns a vector v, which in turn, is dotted with the input vector costs.

def expensive_computation(x : jax.Array, key : jax.random.PRNGKey) -> jax.Array:
     Does some operations that are slow...
     return output

@jax.custom_jvp
def forward(x : jax.Array, costs : jax.Array, key : jax.random.PRNGKey) -> jax.Array:
     v = expensive_computation(x, key)
     return jnp.dot(costs, v)

Lets assume for simplicity that we are only ever going to take the derivative of forward with respect to its first input argument x. We can implement a custom jvp rule for forward using the notation from the Jax docs online (for illustrative purposes, lets assume that the custom jvp is the dot product of v and the tangent vector corresponding to x):

@forward.defjvp
def forward_jvp(primals, tangents):
      x, costs, key = primals
      tangent_x, _ , _ = tangents # only tangent for x is needed.

      # recompute the primal_out
     v = expensive_computation(x, key)
     primals_out = jnp.dot(costs, v)
     
     # compute the tangent_out corresponding to x
     tangent_out_x = jnp.dot(v, tangent_x)
    # no need for calculating tangent_out for costs and key.
     tangent_out = (tangent_out_x, None, None) 
     return primal_out, tangent_out

My questions are the following:

• For forward-mode auto-diff, am I correct in thinking that jax.value_and_grad(forward) results in expensive_computation being called only once?
• What about for reverse-mode?
• What about reverse mode wrapped in jax.jit?

Another way to implement this custom_jvp is the following:

def g(tangent_x, primal_out, x, costs, key): 
       """
       custom_jvp for forward with respect to first variable x
       """
       return jnp.dot(expensive_computation(x, key), tangent_x)

forward.def_jvps(
      g,
      None,
      None
)
   

This differs from the previous custom_jvp in the way that g only has access to the primal_out and not v. In this case:

• For the above am I correct in assuming that for forward mode differentiation (with no jax.jit wrapper), when computing jax.value_and_grad(forward), the function expensive_computation will be called twice ?
• What about if a jax.jit wrapper is used (would XLA be able to understand that both g and forward share a computation step)?
• Finally, what about for reverse mode (with and without jax.jit)?

Many thanks for your help. Please do let me know if any part of the question is unclear.

[jakevdp]

If you want to figure out exactly which computations are being performed in a given jax operations, I'd suggest using jax.make_jaxpr (see https://jax.readthedocs.io/en/latest/jaxpr.html for details). If I were trying to answer this question for you, the first thing I'd do is turn your pseudocode into real code, putting something like lax.sin in place of expensive_operation, then use jax.make_jaxpr to see how many times lax.sin is called.

But there's no reason for me to do that: it's easy enough for you to do yourself! Hope that helps.

[LawrenceMMStewart]

Hi

Thanks for the response and suggesting jaxprs (I had never used them before)! I did what you suggested and got the following results (displayed below). Indeed, sin is called twice in the second implementation.

I still have a couple of questions regarding the above (which are not answered using jaxpr). If you have any answers @jakevdp they would be greatly appreciated!

• In the code below, how do I know if jax is using forward or reverse mode autodiff? I ask this as maybe sin is called once for forward mode but twice for backward mode (I could see a scenario where v could require recomputation in backward mode).
• When using jax.jit, the generated jaxpr for the second method (i.e. the one that calls sin twice) will be given to the XLA compiler. Is there any way of knowing if the compiler would realise sin(x) is called twice and optimize?

Code to generate jaxpr:

import jax
import numpy as np
import jax.numpy as jnp


@jax.custom_jvp
def f1(x, costs):
    v = jnp.sin(x)
    return jnp.dot(v, x)

@f1.defjvp
def f1jvp(primals, tangents):
    x, costs= primals
    tangent_x, _ = tangents # only tangent for x is needed.

    # recompute the primal_out
    v = jnp.sin(x)
    primal_out = jnp.dot(v, x)

    # compute the tangent_out corresponding to x
    tangent_out= jnp.dot(v, tangent_x)

    # no need for calculating tangent_out for costs and key.
    return primal_out, tangent_out


@jax.custom_jvp
def f2(x, costs):
    v = jnp.sin(x)
    return jnp.dot(v, x)

def g(tangent_x, primal_out, x, costs):
    v = jnp.sin(x)
    return jnp.dot(v, tangent_x)

f2.defjvps(g, None)



x = jax.random.normal(jax.random.PRNGKey(0), (3, ))
costs = jax.random.normal(jax.random.PRNGKey(1), (3, ))


e1 = jax.make_jaxpr(jax.value_and_grad(f1))
e2 = jax.make_jaxpr(jax.value_and_grad(f2))

print("First implementation without jit:")
print(e1(x, costs), '\n')
print("Second implementation without jit:")
print(e2(x, costs), '\n')


e1j = jax.make_jaxpr(jax.jit(jax.value_and_grad(f1)))
e2j = jax.make_jaxpr(jax.jit(jax.value_and_grad(f2)))
print("First implementation with jit:")
print(e1j(x, costs))
print("Second implementation with jit:")
print(e2j(x, costs))

Generated jaxprs:

First implementation without jit:
{ lambda ; a:f32[3] b:f32[3]. let
    _:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] 0.0
    c:f32[3] = sin a
    d:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] c a
    e:f32[3] = dot_general[dimension_numbers=(([], []), ([], []))] 1.0 c
  in (d, e) }

Second implementation without jit:
{ lambda ; a:f32[3] b:f32[3]. let
    _:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] 0.0
    c:f32[] = custom_jvp_call[
      call_jaxpr={ lambda ; d:f32[3] e:f32[3]. let
          f:f32[3] = sin d
          g:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] f d
        in (g,) }
      jvp_jaxpr_thunk=<function _memoize.<locals>.memoized at 0x12312b310>
      num_consts=0
      symbolic_zeros=False
    ] a b
    h:f32[3] = sin a
    i:f32[3] = dot_general[dimension_numbers=(([], []), ([], []))] 1.0 h
  in (c, i) }

First implementation with jit:
{ lambda ; a:f32[3] b:f32[3]. let
    c:f32[] d:f32[3] = pjit[
      jaxpr={ lambda ; e:f32[3] f:f32[3]. let
          _:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] 0.0
          g:f32[3] = sin e
          h:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] g e
          i:f32[3] = dot_general[dimension_numbers=(([], []), ([], []))] 1.0 g
        in (h, i) }
      name=f1
    ] a b
  in (c, d) }

Second implementation with jit:
{ lambda ; a:f32[3] b:f32[3]. let
    c:f32[] d:f32[3] = pjit[
      jaxpr={ lambda ; e:f32[3] f:f32[3]. let
          _:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] 0.0
          g:f32[] = custom_jvp_call[
            call_jaxpr={ lambda ; h:f32[3] i:f32[3]. let
                j:f32[3] = sin h
                k:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] j
                  h
              in (k,) }
            jvp_jaxpr_thunk=<function _memoize.<locals>.memoized at 0x123130670>
            num_consts=0
            symbolic_zeros=False
          ] e f
          l:f32[3] = sin e
          m:f32[3] = dot_general[dimension_numbers=(([], []), ([], []))] 1.0 l
        in (g, m) }
      name=f2
    ] a b
  in (c, d) }

[jakevdp]

In the code below, how do I know if jax is using forward or reverse mode autodiff?

value_and_grad always uses reverse-mode autodiff.

When using jax.jit, the generated jaxpr for the second method (i.e. the one that calls sin twice) will be given to the XLA compiler. Is there any way of knowing if the compiler would realise sin(x) is called twice and optimize?

You can use ahead-of-time compilation and print the compiled HLO; see https://jax.readthedocs.io/en/latest/aot.html. For your code, printing the compiled HLO would look like this:

print(jax.jit(jax.value_and_grad(f1)).lower(x, costs).compile().as_text())`

[LawrenceMMStewart]

You can use ahead-of-time compilation and print the compiled HLO; see https://jax.readthedocs.io/en/latest/aot.html

Thanks, using the link you shared I managed to generate the following graphs when calling value_and_grad on both f1 and f2 (and using jax.jit):

For f1:

For f2:

As we can see in both changes the sin is being called twice.

value_and_grad always uses reverse-mode autodiff.

I believe that this is the reason that sin is being called twice. If one inspects the HLO graph for f1 in forward mode diff, the sin is only called once. To avoid this problem I assume it is essential to define a custom_vjp for f1 in addition to the above custom_jvp, which would avoid calling sin for a second time on the backward.

I will implement the above in the coming days and post the results here. If that is indeed the case I will close the discussion.

@jakevdp I cant seem to find in the docs how Jax chooses between reverse-mode or forward-mode diff (or alternatively which functions use forward/ reverse mode). Am i correct in thinking that only jax.jacfwd and jax.jvp execute in forward mode, and everything else is reverse mode?

Many thanks for your help!

[jakevdp]

jax.jvp is forward-mode autodiff, and jax.vjp is reverse-mode autodiff. As the user, you choose whether to use forward-mode or reverse-mode by choosing which autodiff transformation to use. For example, jax.jacrev, jax.jacobian, and jax.grad rely on jax.vjp, so they use reverse-mode. jax.jacfwd computes the jacobian using forward-mode. Does that make sense?

[froystig]

One reason you're seeing redundant computation is due to the use of the defjvps method (plural) rather than defjvp. The former reduces to the latter after calling the original primal function to compute the primal output. If you'd like to avoid the common subexpressions, use defjvp.

The considerations of JVP versus VJP are separate. It is often a good idea to customize JVPs when possible, so that your code supports forward-mode AD. I find that it is often easier to write JVPs as well. Still, you can avoid these common subexpressions in either setting.

[LawrenceMMStewart]

Hi @froystig, hope you are getting on well and thanks for the clarification. However, I am still a little confused by some of the behaviour I have seen. I have provided a minimal example below demonstrating everything in your comment:

One reason you're seeing redundant computation is due to the use of the defjvps method (plural) rather than defjvp. The former reduces to the latter after calling the original primal function to compute the primal output. If you'd like to avoid the common subexpressions, use defjvp.

(which I fully agree with) and some additional behaviour which I find confusing.

Lets consider the functions f1 and f2 from the above dicussion with some dummy variables to call them on:

# some dummy variables
x = jnp.ones((3,)) * 2.0
costs = jnp.arange(3) * 0.5

Recall that f1 follows the methodology you suggested in your comment, and f2 does not (so we would expect redundant calculations to be made).

Lets first look at the jaxpr that is created when we call f1. Calling jax.make_jaxpr(f1)(x, costs) returns:

{ lambda ; a:f32[3] b:f32[3]. let
    c:f32[] = custom_jvp_call[
      call_jaxpr={ lambda ; d:f32[3] e:f32[3]. let
          f:f32[3] = sin d
          g:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] f d
        in (g,) }
      jvp_jaxpr_thunk=<function _memoize.<locals>.memoized at 0x11f44b430>
      num_consts=0
      symbolic_zeros=False
    ] a b
  in (c,) }

This makes sense, as we are only observing calculations from the forward pass! If we were to also print jax.make_jaxpr(f2)(x, costs), we would see a similar jaxpr (which I will omit to not clutter the example).

As you pointed out jax.value_and_grad(f1), should only calculate sin one time (since it uses custom_jvp not custom_jvps). We can verify this by calling jax.make_jaxpr(jax.value_and_grad(f1))(x, costs):

{ lambda ; a:f32[3] b:f32[3]. let
    _:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] 0.0
    c:f32[3] = sin a
    d:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] c a
    e:f32[3] = dot_general[dimension_numbers=(([], []), ([], []))] 1.0 c
  in (d, e) }

Even though jax.grad is reverse-mode, jax is transforming the vjp into a jvp which will be followed by transposition, which nicely avoids sin being called twice. (@froystig please let me know if I have misunderstood this).

As we agree in our above comments, calling jax.value_grad(f2) will result in redundant computation because of the custom_jvps definition . More precisely,sin will be called once for the forward pass, and a second time when invoking the custom_jvp corresponding the the input variable x. Printing the jaxpr jax.make_jaxpr(jax.value_grad(f2))(x, costs) verifies this:

{ lambda ; a:f32[3] b:f32[3]. let
    _:f32[3] = broadcast_in_dim[broadcast_dimensions=() shape=(3,)] 0.0
    c:f32[] = custom_jvp_call[
      call_jaxpr={ lambda ; d:f32[3] e:f32[3]. let
          f:f32[3] = sin d
          g:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] f d
        in (g,) }
      jvp_jaxpr_thunk=<function _memoize.<locals>.memoized at 0x12e827d30>
      num_consts=0
      symbolic_zeros=False
    ] a b
    h:f32[3] = sin a
    i:f32[3] = dot_general[dimension_numbers=(([], []), ([], []))] 1.0 h
  in (c, i) }

Indeed sin is called once in the custom_jvp_call on variable d and once again on variable a.

So far so good!

However my understanding breaks down when using jax.jit. As suggested by @jakevdp, we can visualize the HLO computations.

Lets visualize the HLO graph when calling jitted versions of f1, jax.value_and_grad(f1) jax.grad(f1).

For jax.jit(f1) we can run print(jax.jit(f1).lower(x, costs).compile().as_text()) which returns:

HloModule jit_f1, entry_computation_layout={(f32[3]{0}, f32[3]{0})->f32[]}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation (param_0.1: f32[3]) -> f32[] {
  %param_0.1 = f32[3]{0} parameter(0)
  %sine.0 = f32[3]{0} sine(f32[3]{0} %param_0.1), metadata={op_name="jit(f1)/jit(main)/sin" source_file="/var/folders/6q/_qrktmj50h3bmp0bz2_sz9_c0000gn/T/ipykernel_17834/3435884880.py" source_line=8}
  ROOT %dot.0 = f32[] dot(f32[3]{0} %sine.0, f32[3]{0} %param_0.1), lhs_contracting_dims={0}, rhs_contracting_dims={0}, metadata={op_name="jit(f1)/jit(main)/dot_general[dimension_numbers=(((0,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/var/folders/6q/_qrktmj50h3bmp0bz2_sz9_c0000gn/T/ipykernel_17834/3435884880.py" source_line=9}
}

ENTRY %main.5 (Arg_0.1: f32[3], Arg_1.2: f32[3]) -> f32[] {
  %Arg_1.2 = f32[3]{0} parameter(1), sharding={replicated}
  %Arg_0.1 = f32[3]{0} parameter(0), sharding={replicated}
  ROOT %fusion = f32[] fusion(f32[3]{0} %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(f1)/jit(main)/dot_general[dimension_numbers=(((0,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/var/folders/6q/_qrktmj50h3bmp0bz2_sz9_c0000gn/T/ipykernel_17834/3435884880.py" source_line=9}
}

We can see that a fused computation has been created which calculates basically calculates f(x) = dot(sin(x), x) for the specific x that is given in the input of f1, i.e. what it is referring to as param_0.1. From my understanding, the call just sets ROOT to be equal to this value (which is what we would expect since that is the entire forward pass). Please let me know if I have misunderstood @froystig .

Now if we look at jax.grad(f1) by running print(jax.jit(jax.value_and_grad(f1)).lower(x, costs).compile().as_text()), we observe:

HloModule jit_f1, entry_computation_layout={(f32[3]{0})->(f32[], f32[3]{0})}, allow_spmd_sharding_propagation_to_output={true,true}

%fused_computation (param_0.1: f32[3]) -> f32[] {
  %param_0.1 = f32[3]{0} parameter(0)
  %sine.0 = f32[3]{0} sine(f32[3]{0} %param_0.1), metadata={op_name="jit(f1)/jit(main)/sin" source_file="/var/folders/6q/_qrktmj50h3bmp0bz2_sz9_c0000gn/T/ipykernel_17834/932336511.py" source_line=17}
  ROOT %dot.2 = f32[] dot(f32[3]{0} %sine.0, f32[3]{0} %param_0.1), lhs_contracting_dims={0}, rhs_contracting_dims={0}, metadata={op_name="jit(f1)/jit(main)/dot_general[dimension_numbers=(((0,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/var/folders/6q/_qrktmj50h3bmp0bz2_sz9_c0000gn/T/ipykernel_17834/932336511.py" source_line=18}
}

ENTRY %main.7 (Arg_0.1: f32[3]) -> (f32[], f32[3]) {
  %Arg_0.1 = f32[3]{0} parameter(0), sharding={replicated}
  %fusion = f32[] fusion(f32[3]{0} %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(f1)/jit(main)/dot_general[dimension_numbers=(((0,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/var/folders/6q/_qrktmj50h3bmp0bz2_sz9_c0000gn/T/ipykernel_17834/932336511.py" source_line=18}
  %sine.3 = f32[3]{0} sine(f32[3]{0} %Arg_0.1), metadata={op_name="jit(f1)/jit(main)/sin" source_file="/var/folders/6q/_qrktmj50h3bmp0bz2_sz9_c0000gn/T/ipykernel_17834/932336511.py" source_line=17}
  ROOT %tuple.6 = (f32[], f32[3]{0}) tuple(f32[] %fusion, f32[3]{0} %sine.3)
}

From my understanding jit is creating the same fusion_computation as that of jax.jit(f1) i.e. dot(sin(x), x). However, in %sine.3 there is a second call to sin on the same argument %Arg_0.1. This seems very strange as now we have a repeat redundant calculation. Have I correctly understood this @froystig ? If the sin call was instead a call to another function which is very expensive, this would certainly not be desirable (and was the motivation behind the original question).

Can this be avoided by defining a custom_vjp that more explicitly shows how to reuse the calculation of sin(x) in both the primal_out and tangent_out, so that jit can maybe not repeat the calculation?

Many thanks for your help.

[LawrenceMMStewart]

Hi @froystig , @jakevdp,

Just coming back to this discussion to clarify my issue in the simplest possible example. I have tried to make this comment self-contained, so there is no need to refer to the above conversation.

Image we have a function f, for which we wish to define a custom gradient using either custom_vjp or custom_jvp. We would like to call jax.jit(jax.value_and_grad(f)) to obtain the forward and backward of f.

Imagine that part of the computation of the gradient of f involves an expensive calculation, which will already be calculated when calling the forward pass of f.

Goal : We would like to implement the custom gradient in Jax so that the variable corresponding to the expensive calculation is saved for the backward, so we avoid repeating the expensive calculation when calling jax.value_and_grad(f).

Lets consider the following toy example, where as @jakevdp suggested we represent the expensive calculation by jnp.sin. To reiterate, we want to avoid repeat calls jnp.sin!

For clarity, I have provided an implementation using both custom_jvp and custom_vjp (even though their outputs are identical, and the custom_jvp will be converted to a vjp as jax.value_and_grad is reverse-mode):

import jax
import numpy as np
import jax.numpy as jnp
from jaxlib import xla_client

@jax.custom_jvp
def fjvp(x):
    s = jnp.sin(x)
    return jnp.dot(s, x)


@fjvp.defjvp
def fjvp_jvp(primals, tangents):
    x= primals[0]
    tangent_x = tangents[0] # only tangent for x is needed.

    # recompute the primal_out
    s = jnp.sin(x)
    primal_out = jnp.dot(s, x)

    # compute the tangent_out corresponding to x
    c = jnp.cos(x)
    J = s + x * c
    tangent_out= jnp.dot(J, tangent_x)

    return primal_out, tangent_out


@jax.custom_vjp
def fvjp(x):
    s = jnp.sin(x)
    return jnp.dot(s, x)

def fvjp_fwd(x):
    s = jnp.sin(x)
    sdot = jnp.dot(s, x)
    return sdot, (x, s)

def fvjp_bwd(res, g):
    x, s = res
    c = jnp.cos(x)
    J = s + x * c
    return (g * J.T, )


fvjp.defvjp(fvjp_fwd, fvjp_bwd)

If we look at the jaxpr for hjvp = jax.value_and_grad(fjvp) and hvjp = jax.value_and_grad(fvjp) we obtain exactly what we would like to see:

"jaxpr for hjvp = jax.value_and_grad(jvp)"
 { lambda ; a:f32[3]. let
    b:f32[3] = sin a
    c:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] b a
    d:f32[3] = cos a
    e:f32[3] = mul a d
    f:f32[3] = add b e
    g:f32[3] = dot_general[dimension_numbers=(([], []), ([], []))] 1.0 f
  in (c, g) }

  "jaxpr for hvjp = jax.value_and_grad(fvjp)"
 { lambda ; a:f32[3]. let
    b:f32[3] = sin a
    c:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] b a
    d:f32[3] = cos a
    e:f32[3] = mul a d
    f:f32[3] = add b e
    g:f32[3] = mul 1.0 f
  in (c, g) }

Just as we wanted, the sin operation is being saved, (for use when calculating c and f) so we have avoided recalculating b.

However, the problem occurs when calling jax.jit of both hjvp and hvjp.

Using .lower() and .compile() we can view the HLO computation graph for both jax.jit(hvjp) and jax.jit(hvjp):

Below are png files depicting the HLO computation graph:

HLO graph of jax.jit(jax.value_and_grad(fjvp):

HLO graph of jax.jit(jax.value_and_grad(fvjp):

As we can see, when performing optimizations the XLA compiler leads to repeat calls of sin. Whilst this does not matter for the toy example, I am seeing the same behaviour for our research project, where the expensive computation contains a computation heavy lax.while_loop, which we certainly do not want to call twice.

Is there some way of structuring my jax code to avoid the recomputation occurring when calling jax.jit, or even forcing the XLA compiler to change its HLO graph?

Many thanks for your help.
Lawrence

[jakevdp]

I suspect in this case, the compiler noticed that sin is actually not that expensive a computation, and decided that it would be optimal to not save the intermediate value for use in both fusions.

[LawrenceMMStewart]

Verified with the following example:

@jax.custom_jvp
def f(A):
    eA = jnp.expm1(A)
    return jnp.sum(eA)

@f.defjvp
def fjvp(primals, tangents):
    A = primals[0]
    tangent_A = tangents[0]

    eA = jnp.expm1(A)

    primal_out = jnp.sum(eA)
    tangent_out = jnp.dot(eA[0,0], tangent_A[0, 0])
    return primal_out, tangent_out

As we can see below this behaviour changes when sin is a more expensive computation.

Thanks for your time and help, closing the discussion now !