Generic support for JIT compilation with custom NumPy ops

It would be great to be able to jit functions that make use of custom CPU operation, i.e., implemented with NumPy arrays. This would be a really valuable extension point for integrating JAX with existing codes/algorithms, and would possibly solve the final remaining use-cases for autograd.

Right now, you can use custom CPU operations if you don’t jit, but that adds a very large amount of dispatch overhead.

My understanding is this could be possible by making use of XLA’s CustomCall support.

[shoyer]

I think we can consider this fixed by the new (experimental) version of host_callback.call, e.g.,

import jax.experimental.host_callback as hcb
import jax.numpy as jnp
import jax

def myprint(x):
  print('inside myprint:', type(x), x)
  return x

@jax.jit
def device_fun(x):
  return hcb.call(myprint, x, result_shape=x)

device_fun(jnp.arange(10))

Prints: inside myprint: <class 'numpy.ndarray'> [0 1 2 3 4 5 6 7 8 9]

Please give this a try and file new issues CCing @gnecula if you encounter any issues!

[Thenerdstation]

Bump on this.

We would like to be able to call python code like how tf.py_func works. I understand that type/shape inference is one of the blockers here, but having some kind of “manual” typing would work for us.

So something like

def f(a: np.array) -> np.array:
  res = ...# some non-jax math
  return res

typed_f = jax.give_signature(
  function=f,
  input=Shape((1,), jnp.float32), 
  output=Shape((1,), jnp.float32))

@jax.jit
def g(a)
  return typed_f(a) # Should work assuming our earlier code worked