bijx.Distribution

class bijx.Distribution

Bases: Module

Base class for all probability distributions.

Provides the fundamental interface for sampling and density evaluation that all distribution implementations must follow. Supports both explicit random key passing and automatic key management through the rngs attribute.

The base class does not force a single event shape to be stored, which would be incompatible with general pytree objects that have different types of leaves. Instead, the get_batch_shape method must be implemented to extract the batch shape given a (batched or not) sample.

The child class ArrayDistribution should be used for simple distributions over arrays (not pytree objects).

Parameters:

rngs (Rngs | None) – Optional random number generator state for automatic key management.

Note

Subclasses must implement get_batch_shape(), sample(), and log_density(). The density() method is automatically derived from log_density().

__init__(rngs=None)

Parameters:

rngs (Rngs | None)

Methods

density(x, **kwargs)	Evaluate probability density at given points.
get_batch_shape(x)	Extract batch dimensions from a sample.
log_density(x, **kwargs)	Evaluate log probability density at given points.
sample([batch_shape, rng])	Generate samples from the distribution.

get_batch_shape(x)

Extract batch dimensions from a sample.

Parameters:

x (Any) – A sample from this distribution.

Return type:

tuple[int, ...]

Returns:

Tuple representing the batch dimensions of the sample.

sample(batch_shape=(), rng=None, **kwargs)

Generate samples from the distribution.

Parameters:

• batch_shape (tuple[int, ...]) – Shape of batch dimensions for vectorized sampling.

• rng (Array | None) – Random key for sampling, or None to use internal rngs.

• **kwargs – Additional distribution-specific sampling arguments.

Return type:

tuple[Any, Array]

Returns:

Tuple of (samples, log_densities) where samples have shape (*batch_shape, *event_shape) and log_densities have shape batch_shape.

log_density(x, **kwargs)

Evaluate log probability density at given points.

Parameters:

• x (Any) – Points at which to evaluate density, with event dimensions matching the distribution’s event shape.

• **kwargs – Additional distribution-specific evaluation arguments.

Return type:

Array

Returns:

Log density values with batch dimensions matching input.

density(x, **kwargs)

Evaluate probability density at given points.

Parameters:

• x (Any) – Point(s) at which to evaluate density.

• **kwargs – Additional distribution-specific evaluation arguments.

Return type:

Array

Returns:

Density values (exponential of log density).