Support for automatic differentiation of JAX code

Hi,

I was wondering if using enzyme to automatically derive the VJP/JVP of JAX functions rather than only C++ code is possible with this project?

As a simple motivating example, there are cases where using reverse-mode AD results in poor performance of functions that are composed of many small ops as propagating derivatives through each intermediate node in the tape requires at least one read and write. I have given a simple example below where a hand-derived VJP results in >2x speedup. It would be very useful if Enzyme could be used to automatically derive such custom JVP/VJPs for arbitrary JAX code in cases where regular tape-based reverse-mode AD yields slow code.

I have only a surface level understanding of compilers, IR and enzyme, but it seems to me from skimming through the code in this repo that the pieces to support this functionality are already in place, and that getting the StableHLO MLIR from compiled functions in JAX is straight-forward. Is an API for this functionality already planned, or would you be interested in adding it?

```python
import time

import jax
import jax.numpy as jnp
from jax import custom_vjp

def f_standard(x, a):
return jnp.sin(a * x) ** 2 / a

@custom_vjp
def f_custom(x, a):
return jnp.sin(a * x) ** 2 / a

def f_custom_bwd(res, g):
x, a = res
ax = a * x
sin_ax = jnp.sin(ax)
sin_2ax = jnp.sin(2 * ax)
grad_x = g * sin_2ax
grad_a_local = g * (ax * sin_2ax - sin_ax**2) / (a**2)
grad_a = grad_a_local.sum(axis=tuple(range(grad_a_local.ndim - 1)))
return grad_x, grad_a

f_custom.defvjp(lambda x, a: (f_custom(x, a), (x, a)), f_custom_bwd)

def bench(fn, x, a, n=100):
@jax.jit
def vjp_fn(x, a, g):
_, vjp = jax.vjp(fn, x, a)
return vjp(g)

g = jnp.ones_like(x)
for _ in range(10):
vjp_fn(x, a, g)[0].block_until_ready()
start = time.perf_counter()
for _ in range(n):
vjp_fn(x, a, g)[0].block_until_ready()
return (time.perf_counter() - start) / n * 1000

key = jax.random.PRNGKey(0)
B, C = 32, 64

x = jax.random.normal(key, (B, 512, C))
a = jax.random.uniform(key, (C,), minval=0.5, maxval=2.0)
print(jax.jit(f_standard).lower(x, a).as_text())

# Correctness check
g = jnp.ones_like(x)
(gx_std, ga_std) = jax.vjp(f_standard, x, a)[1](g)
(gx_cst, ga_cst) = jax.vjp(f_custom, x, a)[1](g)
assert jnp.allclose(gx_std, gx_cst), f"grad_x mismatch: {jnp.abs(gx_std - gx_cst).max()}"
assert jnp.allclose(ga_std, ga_cst), f"grad_a mismatch: {jnp.abs(ga_std - ga_cst).max()}"
print("Correctness check passed")

print(f"{'(B,T,C)':<20} {'Std (ms)':>10} {'Cust (ms)':>10} {'Speedup':>10}")
for T in [512, 1024, 2048, 4096, 8192, 16384]:
x = jax.random.normal(key, (B, T, C))
a = jax.random.uniform(key, (C,), minval=0.5, maxval=2.0)
t_std = bench(f_standard, x, a)
t_cst = bench(f_custom, x, a)
print(f"{str((B, T, C)):<20} {t_std:>10.3f} {t_cst:>10.3f} {t_std / t_cst:>9.2f}x")

```

On my RTX 4070 Ti:

```
module @jit_f_standard attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<32x512x64xf32>, %arg1: tensor<64xf32>) -> (tensor<32x512x64xf32> {jax.result_info = "result"}) {
%0 = stablehlo.broadcast_in_dim %arg1, dims = [2] : (tensor<64xf32>) -> tensor<1x1x64xf32>
%1 = stablehlo.broadcast_in_dim %0, dims = [0, 1, 2] : (tensor<1x1x64xf32>) -> tensor<32x512x64xf32>
%2 = stablehlo.multiply %1, %arg0 : tensor<32x512x64xf32>
%3 = stablehlo.sine %2 : tensor<32x512x64xf32>
%4 = stablehlo.multiply %3, %3 : tensor<32x512x64xf32>
%5 = stablehlo.broadcast_in_dim %arg1, dims = [2] : (tensor<64xf32>) -> tensor<1x1x64xf32>
%6 = stablehlo.broadcast_in_dim %5, dims = [0, 1, 2] : (tensor<1x1x64xf32>) -> tensor<32x512x64xf32>
%7 = stablehlo.divide %4, %6 : tensor<32x512x64xf32>
return %7 : tensor<32x512x64xf32>
}
}

Correctness check passed
(B,T,C) Std (ms) Cust (ms) Speedup

(32, 512, 64) 0.076 0.060 1.27x
(32, 1024, 64) 0.077 0.061 1.25x
(32, 2048, 64) 0.123 0.085 1.46x
(32, 4096, 64) 0.622 0.299 2.08x
(32, 8192, 64) 1.184 0.543 2.18x
(32, 16384, 64) 2.293 1.021 2.25x
```

Start with the StableHLO MLIR emitted by jax.jit(...).lower(...).as_text() and the proposed Enzyme-JAX integration path. Define whether arbitrary JAX functions can receive generated custom JVP/VJP support, then validate the API with the f_standard/f_custom correctness check and benchmark shown in the issue.

Written by the indexing model from the issue text.