[JAX] Add Python binding for building a colocated Python program

openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators

Apache License 2.0

2.75k stars 440 forks source link

[JAX] Add Python binding for building a colocated Python program #19811

Closed copybara-service[bot] closed 5 days ago

copybara-service[bot] commented 6 days ago

[JAX] Add Python binding for building a colocated Python program

This change adds a Python binding that makes ifrt::CustomCallProgram for a colocated Python program. This Python binding will be used internally in the colocated Python API implementation. The API does not yet compile the program into an executable, which will be added separately.