How to combine jax.experimental.custom_dce with JVP and transpose rules from jax.extend.core.Primitive ?

[mariogeiger]

I have a jax.extend.core.Primitive primitive that implements JVP, transpose, batching.

segmented_polynomial_p.def_abstract_eval(segmented_polynomial_abstract_eval)
segmented_polynomial_p.def_impl(partial(xla.apply_primitive, segmented_polynomial_p))
mlir.register_lowering(
    segmented_polynomial_p,
    mlir.lower_fun(partial(segmented_polynomial_impl, "cuda"), segmented_polynomial_p.multiple_results), "cuda",
)
mlir.register_lowering(
    segmented_polynomial_p,
    mlir.lower_fun(partial(segmented_polynomial_impl, None), segmented_polynomial_p.multiple_results), None,
)
ad.primitive_jvps[segmented_polynomial_p] = segmented_polynomial_jvp
ad.primitive_transposes[segmented_polynomial_p] = segmented_polynomial_transpose
batching.primitive_batchers[segmented_polynomial_p] = segmented_polynomial_batching

I want to add custom_dce to my primitive because I'm using jax.ffi.ffi_call inside partial(segmented_polynomial_impl, "cuda").
I tried to add custom_dce:

1. After AD => inside segmented_polynomial_impl but it gets ignored
2. Before AD => before calling segmented_polynomial_p.bind, it works when I test the DCE but it crash with the following error when I use AD

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

[mariogeiger]

Nevermind, I need to use the underlying partial_eval.dce_rules directly