When are fwd and bwd functions called for custom_vjp? What would cause fwd to run but not bwd?

[tylerflex]

When defining a custom_vjp for a function, the forward and backward functions are generally both run.

Take this simple example:

import jax

@jax.custom_vjp
def f(x):
    return x + 2

def f_fwd(x):
    print('fwd')
    return x + 2, ()

def f_bwd(res, g):
    print('bwd')
    return g,

f.defvjp(f_fwd, f_bwd)
jax.grad(f)(0.0)

fwd
bwd
Array(1., dtype=float32, weak_type=True)

I'm trying to understand under what conditions the f_fwd and f_bwd functions are called.

EDIT: after playing with this a bit more, I realized that I can get f_fwd to be called (without f_bwd) if I write something like

def g(x):
    y = f(x)
    return 1.0

jax.grad(g)(0.0)

How does jax determine whether to omit the call to run_bwd in this case? vs in the case where I return y?