Question about hessians for custom operations from C++

[samuelpmish]

Hi,

I'm interested in using JAX for finite element simulation. Some of the building blocks of finite elements (e.g. interpolation, integration) are expensive linear operators. I have an existing C++/CUDA library with optimized implementations of these operations that I was hoping to use from within JAX (rather than rewrite everything from scratch), but many of the guides I see related to C++ extensions and custom derivatives only describe how to register a VJP and JVP for these custom operations.

Are custom operations from C++ limited to first derivatives, or is there a way to let JAX know that certain operations are linear (so that any higher order derivatives are identically zero)? I love JAX's ability to easily compute hessians, and would like to still be able to take two derivatives of code using the custom C++ interpolation / integration implementations.

Thanks

[jakevdp]

If you want your custom kernels to be compatible with JAX autodiff, you should register custom_vjp and/or custom_jvp rules for them (see Custom derivative rules for JAX-transformable Python functions). As long as the jvp/vjp rules are themselves compatible with autodiff, then you will be order to perform second-order autodiff on these functions.

[jakevdp]

It looks like pure_callback expects a JAX array or tracer as input. You're passing a numpy array. You could either explicitly convert x to jax array with jnp.asarray(), or wrap your function in jit.

[samuelpmish]

I'm getting the same error message after changing it to

def f(x):
  x = jnp.asarray(x)
  result_shape_dtype = jax.ShapeDtypeStruct(shape=(3,), dtype=f64)
  return jax.pure_callback(func, result_shape_dtype, x)

print(f(x))
# > jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: TypeError: func(): incompatible function arguments. The following argument types are supported:
# >    1. func(x: jaxlib.xla_extension.DeviceArray[dtype=float64, shape=(3)]) -> jaxlib.xla_extension.DeviceArray[dtype=float64, shape=(3)]
# >
# > Invoked with types: ndarray

full output

Traceback (most recent call last):
  File "/Users/sam/code/jax_nanobind/jax_test.py", line 37, in <module>
    print(f(x))
          ^^^^
  File "/Users/sam/code/jax_nanobind/jax_test.py", line 21, in f
    return jax.pure_callback(func, result_shape_dtype, x)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/callback.py", line 250, in pure_callback_api
    return pure_callback(callback, result_shape_dtypes, *args,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/callback.py", line 192, in pure_callback
    out_flat = pure_callback_p.bind(
               ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/core.py", line 380, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/core.py", line 383, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/core.py", line 815, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/dispatch.py", line 144, in apply_primitive
    return compiled_fun(*args)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 1349, in __call__
    results = self.xla_executable.execute_sharded(input_bufs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

[jakevdp]

Oh, I see, that error seems to be coming from your C++ function. What input types does it accept? pure_callback always calls the functions with numpy arrays, not JAX arrays.

[jakevdp]

I didn't realize you were trying to call GPU code from JAX. In that case, these docs would be more relevant than pure_callback, which is designed for callbacks to the CPU host: https://jax.readthedocs.io/en/latest/Custom_Operation_for_GPUs.html

[samuelpmish]

that looks promising, thank you for all the help!