References¶

• Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S. and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs. GitHub
• Frostig, R., Johnson, M. J. and Leary, C. (2018). Compiling machine learning programs via high-level tracing. SysML Conference. PDF

Software and specifications:

• PJRT. Integration guide
• StableHLO. Specification
• XLA FFI. Documentation of the custom call interface
• JAX foreign function interface. Tutorial
• JAX-CIP, which uses cyxla to run JAX functions inside SCIP.