Enzyme Dialect
Reactant.MLIR.Dialects.enzyme.addRetvalToTrace Method
addRetvalToTrace
Add the function's return value(s) into the execution trace.
sourceReactant.MLIR.Dialects.enzyme.addSampleToTrace Method
addSampleToTrace
Add a sampled value into the execution trace.
sourceReactant.MLIR.Dialects.enzyme.addSubtrace Method
addSubtrace
Insert a subtrace into a parent trace.
sourceReactant.MLIR.Dialects.enzyme.addWeightToTrace Method
addWeightToTrace
Add the aggregated log-probability weight to the execution trace.
sourceReactant.MLIR.Dialects.enzyme.broadcast Method
broadcast
Broadcast the operand by adding extra dimensions with sizes provided by the shape attribute to the front. For scalar operands, ranked tensor is created.
NOTE: Only works for scalar and ranked tensor operands for now.
sourceReactant.MLIR.Dialects.enzyme.cholesky_solve Method
cholesky_solve
Solves the linear system Ax = b for x using Cholesky decomposition. Assuming A is symmetric positive definite!
sourceReactant.MLIR.Dialects.enzyme.concat Method
concat
Concat list of input arguments into a generic value
sourceReactant.MLIR.Dialects.enzyme.dot Method
dot
Computes the dot product of two 1D tensors (vectors).
sourceReactant.MLIR.Dialects.enzyme.dump Method
dump
Debug operation that dumps a tensor value with a label.
sourceReactant.MLIR.Dialects.enzyme.extract Method
extract
Extract value from batched operand at index
sourceReactant.MLIR.Dialects.enzyme.generate Method
generate
Generate an execution trace and weight from a probabilistic function. If a constraint dict is provided AND the sample op's symbol is in the constrained_symbols array, we will use the corresponding constraint value instead of generating new samples from the probabilistic function. By convention, the 0th operand in inputs or outputs is the initial RNG state (seed).
Reactant.MLIR.Dialects.enzyme.getFlattenedSamplesFromTrace Method
getFlattenedSamplesFromTrace
Get sampled values for multiple addresses from an execution trace and flatten them into a single position vector for HMC.
sourceReactant.MLIR.Dialects.enzyme.getSampleFromConstraint Method
getSampleFromConstraint
Get sampled values from a constraint for a given symbol.
sourceReactant.MLIR.Dialects.enzyme.getSampleFromTrace Method
getSampleFromTrace
Get the sampled value for a given symbol from an execution trace.
sourceReactant.MLIR.Dialects.enzyme.getSubconstraint Method
getSubconstraint
Get a subconstraint from a constraint for a given symbol.
sourceReactant.MLIR.Dialects.enzyme.getSubtrace Method
getSubtrace
Get a subtrace from a trace for a given symbol.
sourceReactant.MLIR.Dialects.enzyme.getWeightFromTrace Method
getWeightFromTrace
Get the accumulated log-probability weight from an execution trace.
sourceReactant.MLIR.Dialects.enzyme.initTrace Method
initTrace
Initialize an execution trace for a probabilistic function.
sourceReactant.MLIR.Dialects.enzyme.loop Method
loop
A counted loop operation that iterates from lowerBound to upperBound by step, carrying iter_args through each iteration. The iteration variable and iter_args are passed to the body region.
Reactant.MLIR.Dialects.enzyme.mcmc Function
mcmc
Perform an MCMC inference step (HMC, NUTS, etc.) on a probabilistic function. This operation proposes a new trace using the specified algorithm, computes the acceptance probability, and returns the updated trace. By convention, the 0th operand in inputs is the initial RNG state and the 0th operand in results is the updated RNG state.
Optional HMC-specific parameters:
mass: Mass matrix (identity assumed if not provided)
step_size: Leapfrong integration step size
num_steps: Number of leapfrog steps
initial_momentum: deterministic initial momentum (debug)
Reactant.MLIR.Dialects.enzyme.mh Method
mh
Perform a Metropolis-Hastings step on a probabilistic function. This operation proposes a new trace by regenerating selected addresses, computes the acceptance probability, and returns the updated trace. By convention, the 0th operand in inputs is the initial RNG state and the 0th operand in results is the updated RNG state.
sourceReactant.MLIR.Dialects.enzyme.random Method
random
Generates random numbers using the rng_distribution algorithm and produces a result tensor.
If rng_distribution = UNIFORM, then the random numbers are generated following the uniform distribution over the interval [a, b). If a >= b, the behavior is undefined.
If rng_distribution = NORMAL, then the random numbers are generated following the normal distribution with mean = a and standard deviation = b. If b < 0, the behavior is undefined.
If rng_distribution = MULTINORMAL, then the random numbers are generated following the multivariate normal distribution with mean = a (scalar or vector) and covariance matrix = b. The parameter b should be a positive definite matrix.
By convention, the 0th operand in inputs is the initial RNG state and the 0th operand in results is the updated RNG state.
sourceReactant.MLIR.Dialects.enzyme.regenerate Method
regenerate
Regenerate selected addresses in a probabilistic function while keeping other addresses fixed to their values in the given trace. By convention, the 0th operand in inputs is the initial RNG state and the 0th operand in results is the updated RNG state.
sourceReactant.MLIR.Dialects.enzyme.sample Method
sample
Sample from a distribution. By convention, the 0th operand in inputs or outputs is the initial RNG state (seed).
Reactant.MLIR.Dialects.enzyme.selectTrace Method
selectTrace
Selects between two !enzyme.Trace values (considered scalars here) based on a tensor<i1> condition.
sourceReactant.MLIR.Dialects.enzyme.simulate Method
simulate
Simulate a probabilistic function to generate execution trace by replacing all SampleOps with distribution calls and recording all sampled values into the trace. This op returns the trace, the weight (accumulated log-probability), and the other outputs. By convention, the 0th operand in inputs or outputs is the initial RNG state (seed).
Reactant.MLIR.Dialects.enzyme.unflatten_slice Method
unflatten_slice
Extract a slice from a 1D position vector starting at the given offset, and reconstruct the original multi-dimensional tensor shape (implied by the type).
sourceReactant.MLIR.Dialects.enzyme.untracedCall Method
untracedCall
Call a probabilistic function without tracing. By convention, the 0th operand in inputs or outputs is the initial RNG state (seed).
Reactant.MLIR.Dialects.enzyme.update Method
update
Update selected addresses in a trace with new values from a position vector, re-evaluate the probabilistic function, and return the updated trace with the new weight (log probability) and updated RNG state. By convention, the 0th operand in inputs is the initial RNG state.
source