Enzyme Dialect

Reactant.MLIR.Dialects.enzyme.addRetvalToTrace Method

addRetvalToTrace

Add the function's return value(s) into the execution trace.

source

Reactant.MLIR.Dialects.enzyme.addSampleToTrace Method

addSampleToTrace

Add a sampled value into the execution trace.

source

Reactant.MLIR.Dialects.enzyme.addSubtrace Method

addSubtrace

Insert a subtrace into a parent trace.

source

Reactant.MLIR.Dialects.enzyme.addTo Method

addTo

TODO

source

Reactant.MLIR.Dialects.enzyme.addWeightToTrace Method

addWeightToTrace

Add the aggregated log-probability weight to the execution trace.

source

Reactant.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.

source

Reactant.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).

source

Reactant.MLIR.Dialects.enzyme.getSampleFromConstraint Method

getSampleFromConstraint

Get sampled values from a constraint for a given symbol.

source

Reactant.MLIR.Dialects.enzyme.getSubconstraint Method

getSubconstraint

Get a subconstraint from a constraint for a given symbol.

source

Reactant.MLIR.Dialects.enzyme.initTrace Method

initTrace

Initialize an execution trace for a probabilistic function.

source

Reactant.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).

source

Reactant.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).

source

Reactant.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).

source