Enzyme main
Loading...
Searching...
No Matches
Public Member Functions | Public Attributes | List of all members
mlir::impulse::HMCContext Struct Reference

#include "MLIR/Interfaces/HMCUtils.h"

Inheritance diagram for mlir::impulse::HMCContext:
Inheritance graph
[legend]
Collaboration diagram for mlir::impulse::HMCContext:
Collaboration graph
[legend]

Public Member Functions

 HMCContext (FlatSymbolRefAttr fn, ArrayRef< Value > fnInputs, ArrayRef< Type > fnResultTypes, Value originalTrace, ArrayAttr selection, ArrayAttr allAddresses, Value invMass, Value massMatrixSqrt, Value stepSize, Value trajectoryLength, int64_t positionSize, ArrayRef< SupportInfo > supports, DictionaryAttr autodiffAttrs={})
 
 HMCContext (FlatSymbolRefAttr logpdfFn, ArrayRef< Value > fnInputs, Value invMass, Value massMatrixSqrt, Value stepSize, Value trajectoryLength, int64_t positionSize, DictionaryAttr autodiffAttrs={})
 
bool hasCustomLogpdf () const
 
int64_t getFullTraceSize () const
 
Type getElementType () const
 
RankedTensorType getPositionType () const
 
RankedTensorType getScalarType () const
 
bool hasConstrainedSupports () const
 
HMCContext withStepSize (Value newStepSize) const
 

Public Attributes

FlatSymbolRefAttr fn
 
ArrayRef< Value > fnInputs
 
SmallVector< Type > fnResultTypes
 
Value originalTrace
 
ArrayAttr selection
 
ArrayAttr allAddresses
 
Value invMass
 
Value massMatrixSqrt
 
Value stepSize
 
Value trajectoryLength
 
int64_t positionSize
 
SmallVector< SupportInfo > supports
 
FlatSymbolRefAttr logpdfFn
 
DictionaryAttr autodiffAttrs
 

Detailed Description

Definition at line 98 of file HMCUtils.h.

Constructor & Destructor Documentation

◆ HMCContext() [1/2]

mlir::impulse::HMCContext::HMCContext ( FlatSymbolRefAttr fn,
ArrayRef< Value > fnInputs,
ArrayRef< Type > fnResultTypes,
Value originalTrace,
ArrayAttr selection,
ArrayAttr allAddresses,
Value invMass,
Value massMatrixSqrt,
Value stepSize,
Value trajectoryLength,
int64_t positionSize,
ArrayRef< SupportInfo > supports,
DictionaryAttr autodiffAttrs = {} )
inline

Definition at line 114 of file HMCUtils.h.

◆ HMCContext() [2/2]

mlir::impulse::HMCContext::HMCContext ( FlatSymbolRefAttr logpdfFn,
ArrayRef< Value > fnInputs,
Value invMass,
Value massMatrixSqrt,
Value stepSize,
Value trajectoryLength,
int64_t positionSize,
DictionaryAttr autodiffAttrs = {} )
inline

Definition at line 129 of file HMCUtils.h.

Member Function Documentation

◆ getElementType()

Type mlir::impulse::HMCContext::getElementType ( ) const
inline

Definition at line 145 of file HMCUtils.h.

References stepSize.

Referenced by mlir::impulse::buildBaseTree(), mlir::impulse::checkTurning(), mlir::impulse::combineTrees(), mlir::impulse::computeIntegrationStep(), mlir::impulse::computePotentialAndGradient(), getPositionType(), getScalarType(), mlir::impulse::InitHMC(), mlir::impulse::SampleHMC(), and mlir::impulse::SampleNUTS().

◆ getFullTraceSize()

int64_t mlir::impulse::HMCContext::getFullTraceSize ( ) const
inline

Definition at line 140 of file HMCUtils.h.

References originalTrace, and traceType().

Referenced by mlir::impulse::computePotentialAndGradient(), mlir::impulse::InitHMC(), and scatterPositionToTrace().

◆ getPositionType()

RankedTensorType mlir::impulse::HMCContext::getPositionType ( ) const
inline

Definition at line 149 of file HMCUtils.h.

References getElementType(), and positionSize.

Referenced by mlir::impulse::buildBaseTree(), mlir::impulse::checkIterativeTurning(), mlir::impulse::checkTurning(), mlir::impulse::combineTrees(), mlir::impulse::computeIntegrationStep(), mlir::impulse::computePotentialAndGradient(), mlir::impulse::getLeafFromTree(), mlir::impulse::InitHMC(), mlir::impulse::SampleHMC(), and mlir::impulse::SampleNUTS().

◆ getScalarType()

RankedTensorType mlir::impulse::HMCContext::getScalarType ( ) const
inline

Definition at line 153 of file HMCUtils.h.

References getElementType().

Referenced by mlir::impulse::buildBaseTree(), mlir::impulse::checkTurning(), mlir::impulse::combineTrees(), mlir::impulse::computeIntegrationStep(), mlir::impulse::computePotentialAndGradient(), mlir::impulse::InitHMC(), mlir::impulse::SampleHMC(), and mlir::impulse::SampleNUTS().

◆ hasConstrainedSupports()

bool mlir::impulse::HMCContext::hasConstrainedSupports ( ) const
inline

Definition at line 157 of file HMCUtils.h.

References supports.

◆ hasCustomLogpdf()

bool mlir::impulse::HMCContext::hasCustomLogpdf ( ) const
inline

Definition at line 138 of file HMCUtils.h.

References logpdfFn.

Referenced by mlir::impulse::computePotentialAndGradient(), mlir::impulse::InitHMC(), mlir::impulse::SampleHMC(), and mlir::impulse::SampleNUTS().

◆ withStepSize()

HMCContext mlir::impulse::HMCContext::withStepSize ( Value newStepSize) const
inline

Definition at line 165 of file HMCUtils.h.

References stepSize.

Referenced by mlir::impulse::SampleHMC().

Member Data Documentation

◆ allAddresses

ArrayAttr mlir::impulse::HMCContext::allAddresses

Definition at line 104 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), and mlir::impulse::InitHMC().

◆ autodiffAttrs

DictionaryAttr mlir::impulse::HMCContext::autodiffAttrs

Definition at line 112 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), and mlir::impulse::InitHMC().

◆ fn

FlatSymbolRefAttr mlir::impulse::HMCContext::fn

Definition at line 99 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), and mlir::impulse::InitHMC().

◆ fnInputs

ArrayRef<Value> mlir::impulse::HMCContext::fnInputs

Definition at line 100 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), and mlir::impulse::InitHMC().

◆ fnResultTypes

SmallVector<Type> mlir::impulse::HMCContext::fnResultTypes

Definition at line 101 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), and mlir::impulse::InitHMC().

◆ invMass

Value mlir::impulse::HMCContext::invMass

Definition at line 105 of file HMCUtils.h.

Referenced by mlir::impulse::buildBaseTree(), mlir::impulse::checkTurning(), mlir::impulse::computeIntegrationStep(), mlir::impulse::SampleHMC(), and mlir::impulse::SampleNUTS().

◆ logpdfFn

FlatSymbolRefAttr mlir::impulse::HMCContext::logpdfFn

Definition at line 111 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), hasCustomLogpdf(), and mlir::impulse::InitHMC().

◆ massMatrixSqrt

Value mlir::impulse::HMCContext::massMatrixSqrt

Definition at line 106 of file HMCUtils.h.

Referenced by mlir::impulse::SampleHMC(), and mlir::impulse::SampleNUTS().

◆ originalTrace

Value mlir::impulse::HMCContext::originalTrace

Definition at line 102 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), getFullTraceSize(), and mlir::impulse::InitHMC().

◆ positionSize

int64_t mlir::impulse::HMCContext::positionSize

Definition at line 109 of file HMCUtils.h.

Referenced by mlir::impulse::buildTree(), mlir::impulse::checkIterativeTurning(), mlir::impulse::computePotentialAndGradient(), gatherPositionFromTrace(), getPositionType(), and mlir::impulse::InitHMC().

◆ selection

ArrayAttr mlir::impulse::HMCContext::selection

Definition at line 103 of file HMCUtils.h.

◆ stepSize

Value mlir::impulse::HMCContext::stepSize

Definition at line 107 of file HMCUtils.h.

Referenced by mlir::impulse::buildTree(), mlir::impulse::computeIntegrationStep(), gatherPositionFromTrace(), getElementType(), mlir::impulse::SampleHMC(), scatterPositionToTrace(), and withStepSize().

◆ supports

SmallVector<SupportInfo> mlir::impulse::HMCContext::supports

Definition at line 110 of file HMCUtils.h.

Referenced by mlir::impulse::computePotentialAndGradient(), gatherPositionFromTrace(), hasConstrainedSupports(), mlir::impulse::InitHMC(), and scatterPositionToTrace().

◆ trajectoryLength

Value mlir::impulse::HMCContext::trajectoryLength

Definition at line 108 of file HMCUtils.h.

Referenced by mlir::impulse::SampleHMC().

The documentation for this struct was generated from the following file:
Generated on Fri May 8 2026 19:56:27 for Enzyme by doxygen 1.12.0