bijx.IndependentNormal

class bijx.IndependentNormal[source]

Bases: ArrayDistribution

Independent standard normal distribution over arrays.

Each element of the array is independently distributed as a standard normal distribution \(\mathcal{N}(0, 1)\). The total log density is the sum of individual element log densities.

Example

>>> dist = IndependentNormal(event_shape=(10,), rngs=rngs)
>>> x, log_p = dist.sample(batch_shape=(5,))
>>> assert x.shape == (5, 10)
>>> assert log_p.shape == (5,)
__init__(event_shape, rngs=None)
Parameters:
  • event_shape (tuple[int, ...])

  • rngs (Rngs | None)

Methods

density(x, **kwargs)

Evaluate probability density at given points.

get_batch_shape(x)

Extract batch dimensions from an array sample.

log_density(x, **kwargs)

Evaluate log probability density at given points.

sample([batch_shape, rng])

Generate samples from the distribution.

Attributes

event_axes

Axis indices corresponding to event dimensions.

event_dim

Number of event dimensions.

event_size

Total number of elements in the event shape.

sample(batch_shape=(), *, rng=None, **kwargs)[source]

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[Array, 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)[source]

Evaluate log probability density at given points.

Parameters:
  • x (Union[Array, 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.