JAX computations often need to interact with other parts of the scientific Python stack or require specialized low-level functionality. This chapter addresses how to bridge JAX with external systems and extend its core operations.

You will learn practical techniques for:

• Converting data between JAX and NumPy arrays.
• Employing the DLPack standard for efficient, zero-copy data sharing with other libraries like PyTorch or TensorFlow.
• Executing external Python functions within JAX programs using host_callback and pure_callback, recognizing their use cases and limitations.
• Understanding JAX primitives as the foundational operations known to the system.
• Building custom primitives by defining their abstract evaluation behavior (shapes, dtypes), implementing backend-specific lowering rules (XLA HLO generation), and specifying custom differentiation rules (JVPs/VJPs) to ensure they integrate fully with JAX's transformation system.