Internal API

Private

These functions are not part of the public API and are subject to change at any time.

Reactant.REDUB_ARGUMENTS_NAME Constant

Reactant.REDUB_ARGUMENTS_NAME

The variable name bound to call_with_reactant's tuple of arguments in its @generated method definition.

This binding can be used to manually reference/destructure call_with_reactants arguments

This is required because user arguments could have a name which clashes with whatever name we choose for our argument. Thus we gensym to create it.

This originates from https://github.com/JuliaLabs/Cassette.jl/blob/c29b237c1ec0deda3a1037ec519eebe216952bfe/src/overdub.jl#L154

source

Reactant.Compiler.codegen_unflatten! Function

codegen_unflatten!

Generate Julia code to wrap the XLA buffers back into the output result datatypes. The name is due to its similarity to the unflatten function in jax.tree_util.register_pytree_node.

source

Reactant.Compiler.codegen_flatten! Function

codegen_flatten!

Generate Julia code to extract the XLA buffers from input arguments. The name is due to its similarity to the flatten function in jax.tree_util.register_pytree_node.

Arguments

• linear_args: A list of arguments to be flattened.

Returns

• flatten_names: A list of Symbols representing the names of the flattened arguments.

• flatten_code: A list of Exprs to extract the XLA buffers from the input arguments.

Note

The linearized arguments do not directly refer to the are the arguments that have been flattened into a single list.

source

Reactant.Compiler.codegen_xla_call Function

codegen_xla_call

Generate Julia code to call the XLA executable.

Arguments

• flatten_names: A list of Symbols representing the names of the flattened linear arguments.

• donated_args_mask: A list of UInt8s representing whether the argument is donated.

• nresults: The number of results to expect.

source