Skip to content

Programs

Lowering, compiling and the records they return.

lower 

lower(function: Callable[..., Any], shapes: Sequence[ShapeDtypeStruct | Array], platform: str | None = None, custom_calls: Sequence[str] = ()) -> LoweredProgram

Lowers a function to StableHLO for a platform, which by default follows JAX.

shapes are the arguments as shapes, or as arrays. Their placement and the default device of JAX choose the platform as lowering_platform describes, unless platform names one. The program records the platform, and compile_program loads the plugin for it.

function is typed Callable[..., Any] because it may take and return any arrays that JAX traces. Unused arguments stay in the signature, since callers pass every input of a program. jax.export refuses custom calls that it knows no compatibility guarantees for, such as the targets of third-party libraries, so custom_calls names the targets that the function may call. The kernels of jnp.linalg need no entry.

The function is traced with the current setting of jax.enable_x64, so 64-bit arguments and values need that setting on, as they do in JAX.

ArrayType

Bases: NamedTuple

The element type and shape of an argument or result of a program.

dtype instance-attribute 

dtype: dtype

Type of the elements, one of the dtypes that the runtime supports.

shape instance-attribute 

shape: tuple[int, ...]

Length of each axis, empty for a scalar.

size property 

size: int

Returns the number of elements.

nbytes property 

nbytes: int

Returns the size of the array in bytes.

LoweredProgram

Bases: NamedTuple

A program as StableHLO.

code instance-attribute 

code: bytes

Portable StableHLO bytecode, which the PJRT plugin compiles.

text instance-attribute 

text: str

The module as text, from which the executable cache key is computed.

input_types class-attribute instance-attribute 

input_types: tuple[ArrayType, ...] | None = None

The type of each argument, or None when it is not known.

platform class-attribute instance-attribute 

platform: str | None = None

The platform that the program was lowered for, or None when it is not known.

compile_program 

compile_program(program: LoweredProgram, disk_cache: bool = True) -> tuple[Executable, CompileStatistics]

Compiles a program, or loads its executable from the disk cache.

The plugin for the platform of the program loads first, unless a plugin is loaded already. A cached executable that the plugin rejects is removed and compiled again. The executable takes the argument types of the program, so Executable.run converts its arguments to them.

CompileStatistics

Bases: NamedTuple

Time and counts of compiling programs.

compile_seconds class-attribute instance-attribute 

compile_seconds: float = 0.0

Seconds spent compiling StableHLO with the PJRT plugin.

load_seconds class-attribute instance-attribute 

load_seconds: float = 0.0

Seconds spent reading executables from the disk cache.

compiled class-attribute instance-attribute 

compiled: int = 0

Programs compiled.

loaded class-attribute instance-attribute 

loaded: int = 0

Programs whose executables came from the disk cache.

add 

add(other: CompileStatistics) -> CompileStatistics

Returns the sum of two statistics.

compile_options cached 

compile_options() -> bytes

Returns the serialized CompileOptionsProto for one replica and one partition.

Executable

A program loaded into the PJRT client.

Attributes:

Name Type Description
num_outputs int

The number of outputs of the program.
output_types tuple[ArrayType, ...]

The ArrayType of each output, as the plugin reports it.
input_types tuple[ArrayType, ...] | None

The ArrayType of each argument, or None when it is not known. compile_program sets it from the lowered program, and run then converts its arguments to these types.

output_bytes property 

output_bytes: tuple[int, ...]

The size in bytes of each output, which fills ProgramSpec.output_bytes.

compile staticmethod 

compile(code: bytes, compile_options: bytes) -> Executable

Compiles StableHLO bytecode with a serialized CompileOptionsProto.

deserialize staticmethod 

deserialize(serialized: bytes) -> Executable

Loads an executable that serialize wrote with the same plugin.

run 

run(arguments: Sequence[ArrayLike], zero_copy: bool = True) -> list[np.ndarray]

Executes the program and returns its outputs with their shapes and dtypes.

It goes through execute, as extension modules do, which makes it useful for tests and interactive checks. When the types of the arguments are known, as for an executable from compile_program, each argument converts to the type that the program takes, and otherwise it keeps its own dtype.

serialize 

serialize() -> bytes

Returns the executable in the plugin's serialized form.