JAX python C 回调

Question

Numba 允许直接在 python 中创建 C 回调 @cfunc装饰器（https://numba.pydata.org/numba-doc/0.42.0/user/cfunc.html）：

@cfunc("float64(float64)") 
def square(x):
    return x**2

澄清一下，生成的函数是纯 C 函数，然后可以直接从 C 代码调用。

JAX 中是否有等效的功能（https://jax.readthedocs.io/en/latest/#）？

我已经搜索了一段时间，但没有找到任何东西。我将不胜感激任何提示。

Answer 1

不，JAX 不提供任何类似于 Numba 的 API

cfunc

。