Reactant.jl

Appearance

0.4k

Trusted by

Scientific Computing

Lux.jl

Machine Learning

Configuration Options ​

Scoped Values ​

Warning

Currently options are scattered in the form of global variables and scoped values. We are in the process of migrating all of them into scoped values.

Reactant.with_config Function
julia
with_config(f; kwargs...)

Run the function f within a dynamic scope such that all uses of the config within this scope will use the provided values.

Extended Help

Configuration Options

Lowering

  • lower_partialsort_to_approx_top_k: Whether to lower partialsort and partialsortperm to Ops.approx_top_k. Note that XLA only supports lowering ApproxTopK for TPUs unless fallback_approx_top_k_lowering is set to true.

  • fallback_approx_top_k_lowering: Whether to lower Ops.approx_top_k to stablehlo.top_k if the XLA backend doesn't support ApproxTopK. Defaults to true.

DotGeneral

  • dot_general_algorithm: Algorithm preset for stablehlo.dot_general. Can be nothing, DotGeneralAlgorithm or DotGeneralAlgorithmPreset. Defaults to DotGeneralAlgorithmPreset.DEFAULT.

  • dot_general_precision: Precision for stablehlo.dot_general. Can be nothing, or PrecisionConfig. Defaults to PrecisionConfig.DEFAULT.

  • convolution_precision: Precision for stablehlo.convolution. Can be nothing, or PrecisionConfig. Defaults to PrecisionConfig.DEFAULT.

Zygote Overlay

  • overlay_zygote_calls: Whether to overlay Zygote.gradient calls with Enzyme.autodiff calls. Defaults to true.
source

Values ​

  • DOT_GENERAL_PRECISION: Controls the precision_config for stablehlo.dot_general.

  • DOT_GENERAL_ALGORITHM: Controls the algorithm for stablehlo.dot_general.

  • LOWER_PARTIALSORT_TO_APPROX_TOP_K: Whether to lower partialsort to Ops.approx_top_k. Note that XLA only supports lowering ApproxTopK for TPUs unless FALLBACK_APPROX_TOP_K_LOWERING is set to true. Defaults to false.

  • FALLBACK_APPROX_TOP_K_LOWERING: Whether to fallback to lowering ApproxTopK to stablehlo.top_k if the XLA backend doesn't support ApproxTopK. Defaults to true.

DotGeneral ​

Reactant.DotGeneralAlgorithmPreset Module
julia
DotGeneralAlgorithmPreset

Controls the precision_config for stablehlo.dot_general. Valid values are:

  • DEFAULT

  • ANY_F8_ANY_F8_F32

  • ANY_F8_ANY_F8_F32_FAST_ACCUM

  • ANY_F8_ANY_F8_ANY

  • ANY_F8_ANY_F8_ANY_FAST_ACCUM

  • F16_F16_F16

  • F16_F16_F32

  • BF16_BF16_BF16

  • BF16_BF16_F32

  • BF16_BF16_F32_X3

  • BF16_BF16_F32_X6

  • BF16_BF16_F32_X9

  • F32_F32_F32

  • F64_F64_F64

The following functions are available:

supported_lhs_eltype(dot_algorithm_preset::DotGeneralAlgorithmPreset.T) supported_rhs_eltype(dot_algorithm_preset::DotGeneralAlgorithmPreset.T) accumulation_eltype(dot_algorithm_preset::DotGeneralAlgorithmPreset.T) supported_output_eltype(dot_algorithm_preset::DotGeneralAlgorithmPreset.T, T1, T2) MLIR.IR.Attribute(dot_algorithm_preset::DotGeneralAlgorithmPreset.T, T1, T2)

source
Reactant.PrecisionConfig Module
julia
PrecisionConfig

Controls the precision_config for stablehlo.dot_general. Valid values are:

  • DEFAULT

  • HIGH

  • HIGHEST

The following functions are available:

MLIR.IR.Attribute(precision::PrecisionConfig.T)

source
Reactant.DotGeneralAlgorithm Type
julia
DotGeneralAlgorithm(
    ::Type{lhsT}, ::Type{rhsT}, ::Type{accumT},
    rhs_component_count::Int, lhs_component_count::Int, num_primitive_operations::Int,
    allow_imprecise_accumulation::Bool
)
DotGeneralAlgorithm{lhsT,rhsT,accumT}(
    lhs_component_count::Int, rhs_component_count::Int, num_primitive_operations::Int,
    allow_imprecise_accumulation::Bool
)

Represents the configuration of the stablehlo.dot_general operation.

Arguments

  • lhsT: The type of the left-hand side operand.

  • rhsT: The type of the right-hand side operand.

  • accumT: The type of the accumulation operand.

  • lhs_component_count: The number of components in the left-hand side operand.

  • rhs_component_count: The number of components in the right-hand side operand.

  • num_primitive_operations: The number of primitive operations in the stablehlo.dot_general operation.

source

Zygote Overlay ​

  • OVERLAY_ZYGOTE_CALLS: Whether to overlay Zygote.gradient calls with Enzyme.autodiff calls.

Environment Variables ​

The following environment variables can be used to configure Reactant.

GPU Configuration ​

  • XLA_REACTANT_GPU_MEM_FRACTION: The fraction of GPU memory to use for XLA. Defaults to 0.75.

  • XLA_REACTANT_GPU_PREALLOCATE: Whether to preallocate GPU memory. Defaults to true.

  • REACTANT_VISIBLE_GPU_DEVICES: A comma-separated list of GPU device IDs to use. Defaults to all visible GPU devices. Preferably use CUDA_VISIBLE_DEVICES instead.

TPU Configuration ​

  • TPU_LIBRARY_PATH: The path to the libtpu.so library. If not provided, we download and use Scratch.jl to save the library.

Distributed Setup ​

  • REACTANT_COORDINATOR_BIND_ADDRESS: The address to bind the coordinator to. If not provided, we try to automatically infer it from the environment.