Operator Registration

General Operator Customization Interface

XLACustomOp

Creating a XLA custom call operator.

CPU Operator Customization with Taichi

clear_taichi_aot_caches([kernels])	Clean the cache of the AOT compiled kernels.
count_taichi_aot_kernels()	Count the number of AOT compiled kernels.

CPU Operator Customization with Numba

CustomOpByNumba

Creating a XLA custom call operator with Numba JIT on CPU backend.

compile_cpu_signature_with_numba(c, func, ...)

Operator Autograd Customization

defjvp(primitive, *jvp_rules)

Define JVP rules for any JAX primitive.