Best practices when implementing custom rules

[ASKabalan]

Hello,

I have a question about custom rules like custom_jvp, custom_transpose , custom_batching.

In this case (if I need for example a gradient, a transpose rule and vmap rule)

from jax import custom_jvp
long_iterative_solve = jnp.linalg.solve # for illustration

@custom_jvp
@custom_transpose
def solve(A, b):
  """find x such that A @ x = b"""
  return long_iterative_solve(A, b)

@solve.def_transpose
def solve_transpose(A, tb):
  return solve(A.T, tb)

@solve.defjvp
def solve_jvp(primals, tangents):
  A, b = primals
  tA, tb = tangents
  x = solve(A, b)
  tx = solve(A, tb - tA @ x)    # automatically transposed for VJP
  return x, tx

I saw this example that has multiple decorators stacked but how does this fit if I want to use a vmap aswell.

Does the vmap rule go after or before the vjp?

Thank you

[dfm]

Good question! The dream is (was?) to generally support mixing and matching these decorators, but (as you've probably worked out!) there are some rough edges with the current APIs, and such customization hasn't been widely exercised. For now, you're probably on the right track, but hopefully this will be easier in the future!

As for your specific question about custom_vmap. You typically want custom_vmap on the inside when it is combined with custom_vjp (see the new custom_vmap docstring for an example), which is the most common use case. I expect that you'll want something similar in your case.

[ASKabalan]

Thank you very much for you answer
This looks very good however in my case I need to be able to do both jvp and vjp

This would be great for my use case #22457
However for now a custom_jvp will do just fine

Few problems with enable_custom_vjp_by_custom_transpose

Here is a very long almost MWE

from jax import core, jacfwd
import jax
import jax.numpy as jnp
from jax._src.lib.mlir.dialects import hlo
from jax.interpreters import mlir
from jax import custom_vjp , custom_jvp

square_prim_p = core.Primitive("multiply_add")  # Create the primitive

def square_prim(x):
  return square_prim_p.bind(x)

def square_impl(x):
  return jnp.power(x, 2)

def square_abstract_eval(x):
  return core.ShapedArray(x.shape, x.dtype)

def square_lowering(ctx , xc):
  return hlo.MulOp(xc, xc).results


@custom_vjp
def square_vjp(x):
  return square_prim(x)

def square_vjp_fwd(x):
  return square_vjp(x), x * 2

def square_vjp_bwd(res, g):
  return g * res , 


@custom_jvp
def square_jvp(x):
  return square_prim(x)

@square_jvp.defjvp
def square_jvp_impl(primals, tangents):
  x, = primals
  x_dot, = tangents
  primals_out = square_jvp(x)
  tangents_out = 2 * x_dot * x
  return primals_out, tangents_out


square_vjp.defvjp(square_vjp_fwd, square_vjp_bwd)

square_prim_p.def_impl(square_impl) 
square_prim_p.def_abstract_eval(square_abstract_eval)
mlir.register_lowering(square_prim_p, square_lowering)

vjp_grad = jax.jit(jax.grad(square_vjp))(jnp.array(3.0)) # works
jvp_grad = jax.jit(jax.grad(square_jvp))(jnp.array(3.0)) # works

print(f"vjp_grad: {vjp_grad}")
print(f"jvp_grad: {jvp_grad}")

def _fn_vjp(x):
  dy , vjp = jax.vjp(square_vjp, x)
  return vjp(dy)[0]

def _fn_jvp(x):
  dy , vjp = jax.vjp(square_jvp, x)
  return vjp(dy)[0]

jax.grad(_fn_jvp)(jnp.array(3.0)) 
jax.grad(_fn_vjp)(jnp.array(3.0)) 


from jax import jvp , vjp , jacfwd , jacrev

jvp(square_jvp, (jnp.array(3.0),) , (jnp.array(1.0),))  #OK
vjp(square_vjp, jnp.array(3.0)) #OK
vjp(square_jvp, jnp.array(3.0)) #OK


try:
    jvp(square_vjp, (jnp.array(3.0),) , (jnp.array(1.0),))  #NOT OK
except Exception as e:
    print(f" -> Error when doing a jvp with custom vjp:  \n ----> {e}")

with jax.enable_custom_vjp_by_custom_transpose():
    print(f" -> value of enable_custom_vjp_by_custom_transpose {jax.enable_custom_vjp_by_custom_transpose.value}") # TRUE
    try:
        jvp(square_vjp, (jnp.array(3.0),) , (jnp.array(1.0),))  # STILL NOT OK
    except Exception as e:
        print(f" -> Error when doing a jvp with custom vjp with context custom_vjp:  \n ----> {e}")


with jax.enable_custom_vjp_by_custom_transpose():
    try:
        jacfwd(square_vjp)(jnp.array(3.0)) # NOT OK
    except Exception as e:
        print(f" -> Error when doing a jacfwd with custom vjp with context custom_vjp:  \n ----> {e}")

jacrev(square_jvp)(jnp.array(3.0)) # OK 

the weidest error is ----> <class 'jax._src.interpreters.batching.BatchTrace'> must override process_custom_transpose to handle custom_transpose_call primitives

Using only custom_jvp in my case

However in my much more complicated use case, I am able to do a vjp and jvp on my custom_jvp function .. but when I do a jacrev I get this

NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'custom_partitioning' not implemented

So I need to implement a custom_transpose
custom_transpose is not documented and I am having a hard time finding a way to implement my own like the others

Do you have any pointers on how to overcome this issue

How would it look like in the end ?

Custom transpose -> custom JVP -> custom vmap?

Thank you very much

[ASKabalan]

@dfm I am closing this disscusion because I think that the answer is not yet implemented in JAX

But in my opinion it is a mix of #22457 and #24726 is exactly what I need in addition to supporting custom_partitionning for example with a @def_partition and @def_infer_sharding