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.

Operator Autograd Customization#

defjvp(primitive, *jvp_rules)

Define JVP rules for any JAX primitive.