Struct NUTSChain 

pub struct NUTSChain<T, B, GTarget>
where
    B: AutodiffBackend,
{
    pub position: Tensor<B, 1>,
    /* private fields */
}

Single-chain state and adaptation for NUTS.

Manages the dynamic trajectory building, dual-averaging adaptation of step size, and current position for one chain.

Fields

position: Tensor<B, 1>

Current position in parameter space.

Implementations

impl<T, B, GTarget> NUTSChain<T, B, GTarget>

where T: Float + ElementConversion + Element + SampleUniform + FromPrimitive, B: AutodiffBackend, GTarget: GradientTarget<T, B> + Sync, StandardNormal: Distribution<T>, StandardUniform: Distribution<T>, Exp1: Distribution<T>,

pub fn new( target: GTarget, initial_position: Vec<T>, target_accept_p: T, ) -> Self

Constructs a new NUTSChain for a single chain with the given initial position.

Parameters

• target: The target distribution implementing GradientTarget.
• initial_position: Initial position vector of length D.
• target_accept_p: Desired average acceptance probability for adaptation.

Returns

An initialized NUTSChain.

pub fn set_seed(self, seed: u64) -> Self

Sets a new random seed for this chain to ensure reproducibility.

Parameters

• seed: Seed value for the chain’s RNG.

Returns

self with the RNG re-seeded.

pub fn run(&mut self, n_collect: usize, n_discard: usize) -> Tensor<B, 2>

Runs the chain for n_collect + n_discard steps, adapting during burn-in and returning collected samples.

Parameters

• n_collect: Number of samples to collect after adaptation.
• n_discard: Number of burn-in steps for adaptation.

Returns

A 2D tensor of shape [n_collect, D] containing collected samples.

pub fn step(&mut self)

Performs one NUTS update step, including tree expansion and adaptation updates.

This method updates self.position and adaptation statistics in-place.

Trait Implementations

impl<T: Clone, B, GTarget: Clone> Clone for NUTSChain<T, B, GTarget>

where B: AutodiffBackend + Clone,

fn clone(&self) -> NUTSChain<T, B, GTarget>

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<T: Debug, B, GTarget: Debug> Debug for NUTSChain<T, B, GTarget>

where B: AutodiffBackend + Debug,

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.