Lower 'python' JAX code to MLIR

[ASKabalan]

Hello JAX team,

is there a way that we can get the 'mlir' or 'differentiation rule' of some JAX code similar to how we get the lowered mlir like so

mlir.lower_fun

Basically I am looking to do this

from jax.interpreters import batching, ad, mlir


def my_primitive_impl(x):
    return x **2

mlir.register_lowering(my_primitive_p,mlir.lower_fun(my_primitive_impl))
ad.primitive_jvps[my_primitive_p] = ad.lower_jvp(my_primitive_impl)
batching.primitive_batchers[my_primitive_p] = batching.lower_batcher(my_primitive_impl)

[dfm]

There are ways to achieve this behavior, but it's not always totally straightforward. Can you say more about what your motivation/use case is for this, and perhaps we can work out the best approach?

[ASKabalan]

Hello Dan,

I have a very specific use case actually.
I am inbetween using shard_map and custom_partionning though I found the latter much more powerfull.
When writing a custom partionning rule for a C++ primitive it is very good
But for a pure JAX it is problematic since I have to redefine the the differentiation rule because under the hood a primitive is created in here https://github.com/google/jax/blob/main/jax/_src/custom_partitioning.py#L218
on the other hand shard_map is differentiable
I would rather use custom_partionning rather that shard_map.

I have created a wrapper that handles both cases (Custom_partitionning for pure JAX and for CUDA primitives)
https://github.com/DifferentiableUniverseInitiative/jaxDecomp/blob/main/jaxdecomp/_src/spmd_ops.py#L120

I write the differentiation rule for my cuda primitives because I have to anyway.
But is there a way I can set the custom_vjp rule to be the same one I get when I do grad of the python code.

Thank you.

[dfm]

I must admit that I don't have a lot of experience with custom_partitioning so I can't help too much with that part. But, I can help with the specific question of how to use the existing jax.vjp implementation within a custom_vjp. We can update the custom_vjp example from the docs to do this:

import jax
import jax.numpy as jnp

def f_ref(x, y):
  return jnp.sin(x) * y

@jax.custom_vjp
def f(x, y):
  return f_ref(x, y)

def f_fwd(x, y):
  return f(x, y), (x, y)  # The residuals should _be_ the primals

def f_bwd(res, g):
  _, vjp_fun = jax.vjp(f_ref, *res)
  return vjp_fun(g)

f.defvjp(f_fwd, f_bwd)

Hope this helps!

[ASKabalan]

Perfect!

Exactly what I want
Works for non linear function from what I gather