Variables

parax.AbstractVariable

Bases: AbstractUnwrappable[Array]

The abstract interface for all model variables.

Derive from this class and override value to implement custom variable unwrapping behaviour.

All parameters in Parax, such as parax.Random, parax.Constrained etc., derive from this class.

Corner Case Note (Math & Dunders): Because this class implements the __jax_array__ protocol and all standard math dunder methods, variables can be used directly in JAX expressions without explicitly calling unwrap(). However, applying any math operation (e.g., var + 1) instantly evaluates the value and returns a standard jax.Array, stripping away the metadata and constraint wrappers.

value abstractmethod property

Returns the underlying, fully computed value of the variable.

parax.Param = AbstractVariable | Inexact[Array, '...'] module-attribute

A type alias representing a JAX parameter.

This includes any Parax variables (like Tagged, Constrained, Derived) as well as standard JAX inexact arrays.

parax.Real

Bases: AbstractVariable, AbstractWrappable[Array]

A plane real variable.

Useful as a placeholder, e.g. for frameworks that only want to allow parax.AbstractVariable instances to be trainable.

Attributes:

Name	Type	Description
raw_value	Param	The raw value used by optimizers and samplers.

parax.Tagged

Bases: AbstractVariable, AbstractAnnotated[dict], AbstractWrappable[Array]

A variable with dictionary metadata.

Represents a simple, trainable variable with a single underlying raw_value and metadata.

Attributes:

Name	Type	Description
raw_value	Param	The raw value used by optimizers and samplers.
metadata	dict	Additional arbitrary metadata.

parax.Fixed(raw_value=None)

Bases: AbstractVariable, AbstractConstant[Param], AbstractWrappable[Array]

A fixed variable.

Implements AbstractConstant for filtering during partitioning.

Attributes:

Name	Type	Description
raw_value	Param	The underlying variable that is being fixed.

Parameters:

Name	Type	Description	Default
raw_value	Param | None	The underlying value to be fixed.	None

Source code in parax/variables.py

def __init__(self, raw_value: Param | None = None):
    """
    Args:
        raw_value: The underlying value to be fixed.
    """
    # Error checking
    if isinstance(raw_value, Fixed):
        raw_value = raw_value.raw_value
    self.raw_value = raw_value    

parax.Derived

Bases: AbstractVariable

A derived variable.

The parameter's value is dynamically derived via an arbitrary callable.

This is ideal for one-way transformations, projections, or normalizations where a strict bijector (with an inverse) is not required or mathematically possible (e.g., applying jax.nn.softmax to raw logits).

Attributes:

Name	Type	Description
fn	Callable	The callable used to transform the raw value.
raw_value	Param	The raw value used by optimizers and samplers.

value property

The derived value.

Returns the raw state transformed by the derivation function.

parax.Transformed(bijector, raw_value)

Bases: AbstractVariable, AbstractWrappable[Array]

A variable transformed by a bijector.

The parameter's value is dynamically derived via a bijective transform.

Note that this simply applies forward/inverse passes during unwrapping, and does NOT apply any special treatment to any other variable types (e.g. parax.Constrained or parax.Random variables).

Attributes:

Name	Type	Description
bijector	AbstractBijector	The bijector used to transform the raw value.
raw_value	Param	The raw value used by optimizers and samplers.

Parameters:

Name	Type	Description	Default
bijector	AbstractBijector	The bijector used to transform the raw value.	required
raw_value	Param	The underlying value to be fixed.	required

Source code in parax/variables.py

def __init__(self, bijector: AbstractBijector, raw_value: Param):
    """
    Args:
        bijector: The bijector used to transform the raw value.
        raw_value: The underlying value to be fixed.
    """
    if isinstance(raw_value, Transformed):
        bijector = Chain([bijector, as_unwrapped(raw_value.bijector)])
        raw_value = raw_value.raw_value

    self.bijector = bijector
    self.raw_value = raw_value

value property

The derived value.

Returns the raw state transformed by the derivation function.

parax.Bounded(bounds, raw_value=None)

Bases: AbstractVariable, AbstractBounded[Array], AbstractWrappable[Array]

A bounded variable.

This simply attaches bounds to an existing variable or an array, and does not apply any bijective constraints. For enforcing constraints on an array, use parax.variables.Constrained.

Attributes:

Name	Type	Description
bounds	tuple[Array, Array]	The parameter bounds.
raw_value	Param	The raw, unconstrained value on the real number line.

Parameters:

Name	Type	Description	Default
bounds	tuple[Array, Array]	The parameter bounds.	required
raw_value	Param | None	The underlying value. Must lie within bounds.	None

Source code in parax/variables.py

def __init__(
    self,
    bounds: tuple[Array, Array],
    raw_value: Param | None = None,
):
    """
    Args:
        bounds: The parameter bounds.
        raw_value: The underlying value. Must lie within `bounds`.
    """
    if raw_value is not None:
        raw_array = jnp.asarray(raw_value)
        lower, upper = jnp.asarray(bounds)

        is_out_of_bounds = jnp.any((raw_array < lower) | (raw_array > upper))
        raw_value = eqx.error_if(
            raw_value,
            is_out_of_bounds,
            "Bounded variable initialized with a value outside of its specified bounds."
        )

    self.bounds = bounds
    self.raw_value = raw_value

parax.Constrained(constraint=None, value=None, *, raw_value=None)

Bases: AbstractVariable, AbstractConstrainable[Array], AbstractWrappable[Array]

A constrained variable.

The constraint is specified via a parax.AbstractConstraint.

The constraint is automatically applied as a bijection mapping during evaluation. Implements the parax.bounds.AbstractBounded interface for integration with bounded optimizers.

Attributes:

Name	Type	Description
constraint	AbstractConstraint	The parameter constraint defining bounds and bijector mappings.
raw_value	Array	The raw, unconstrained value on the real number line.

Parameters:

Name	Type	Description	Default
constraint	AbstractConstraint | None	A Parax constraint. If None, defaults to parax.RealLine (unconstrained).	None
value	Array | None	The desired output (constrained) value. If provided, the internal raw_value is computed dynamically via the constraint's inverse bijector. Mutually exclusive with raw_value.	None
raw_value	Array | None	The unconstrained underlying value. Mutually exclusive with value.	None

Source code in parax/variables.py

def __init__(
    self,
    constraint: AbstractConstraint | None = None,
    value: Array | None = None,
    *,
    raw_value: Array | None = None,
):
    """
    Args:
        constraint: A Parax constraint. If None, defaults to `parax.RealLine` (unconstrained).
        value: The desired output (constrained) value. If provided, the internal 
            `raw_value` is computed dynamically via the constraint's inverse bijector. 
            Mutually exclusive with `raw_value`.
        raw_value: The unconstrained underlying value. Mutually exclusive with `value`.
    """
    # Error checking
    if value is None and raw_value is None:
        raise ValueError("Must provide either `value` or `raw_value`.")
    if value is not None and raw_value is not None:
        raise ValueError("Cannot provide both `value` and `raw_value`.")

    # Array standardization
    if raw_value is not None:
        raw_value = jnp.asarray(raw_value)
        shape = raw_value.shape
    else:
        value = jnp.asarray(value)
        shape = value.shape

    # Constraint and distribution standardization
    if constraint is None:
        constraint = RealLine(shape=shape)

    # Raw value standardization
    if value is not None:
        raw_value = constraint.bijector.inverse(value)
        raw_value = eqx.error_if(
            raw_value,
            jnp.any(jnp.isnan(raw_value)),
            "Constraint violated for variable upon initialization (produced NaNs)."
        )

    self.constraint = constraint
    self.raw_value = raw_value

parax.Random(distribution, constraint=None, value=None, *, raw_value=None)

Bases: AbstractVariable, AbstractProbabilistic[Array], AbstractConstrainable[Array], AbstractWrappable[Array]

A random variable with an optional constraint.

The distribution is specified via a distreqx.distributions.AbstractDistribution. The constraint is specified via a parax.constraint.AbstractConstraint.

The variable implements the parax.probability.AbstractProbabilistic interface to integrate with stochastic samplers and other algorithms.

Attributes:

Name	Type	Description
distribution	AbstractDistribution	The probability distribution of raw_value.
constraint	AbstractConstraint	The constraint that defines the support of distribution. Can be None, in which case this function will attempt to automatically infer the constraint from the distribution's using parax.constraints.get_constraint_for_distribution.
raw_value	Array	The raw un-probabilistic value on the real number line. Can be None, in which case the mean of the distribution is used. If the mean is not supported, an exception is thrown.

Parameters:

Name	Type	Description	Default
raw_value	Array | None	The un-probabilistic raw value.	None
distribution	AbstractDistribution	The probability distribution.	required
constraint	AbstractConstraint | None	The distribution's constraint. If None, then a	None

Source code in parax/variables.py

def __init__(
    self,
    distribution: AbstractDistribution,
    constraint: AbstractConstraint | None = None,
    value: Array | None = None,
    *,
    raw_value: Array | None = None,
):
    """
    Args:
        raw_value: The un-probabilistic raw value.
        distribution: The probability distribution.
        constraint: The distribution's constraint. If `None`, then a 
    """
    if value is not None and raw_value is not None:
        raise ValueError("Cannot provide both `value` and `raw_value`.")

    # Derive physical value if both are missing
    if value is None and raw_value is None:
        try:
            value = jnp.array(distribution.mean())
        except Exception:
            raise ValueError(
                "`value` or `raw_value` must be provided if the "
                "distribution does not support `mean()`."
            )

    # Constraint resolution
    if constraint is None:
        constraint = infer_distribution_constraint(distribution)

    # Calculate unconstrained raw_value
    if value is not None:
        raw_value = constraint.bijector.inverse(jnp.asarray(value))
        raw_value = eqx.error_if(
            raw_value,
            jnp.any(jnp.isnan(raw_value)),
            "Constraint violated for variable upon initialization (produced NaNs)."
        )
    else:
        raw_value = jnp.asarray(raw_value)

    self.distribution = distribution
    self.constraint = constraint
    self.raw_value = raw_value

parax.as_param(value)

Returns value as a parax.Param, wrapping it if necessary.

Parameters:

Name	Type	Description	Default
value	Any	An arbitrary value or array.	required

Returns:

Type	Description
Any	The instantiated parameter.

Source code in parax/variables.py

def as_param(value: Any) -> Any:
    """
    Returns `value` as a `parax.Param`, wrapping it if necessary.

    Args:
        value: An arbitrary value or array.

    Returns:
        The instantiated parameter.
    """    
    if is_param(value):
        return value
    return jnp.asarray(value, dtype=float)

parax.as_variable(value)

Returns value as a parax.AbstractVariable, wrapping it if necessary.

Parameters:

Name	Type	Description	Default
value	Any	An arbitrary value or array.	required

Returns:

Type	Description
Any	The instantiated parameter.

Source code in parax/variables.py

def as_variable(value: Any) -> Any:
    """
    Returns `value` as a `parax.AbstractVariable`, wrapping it if necessary.

    Args:
        value: An arbitrary value or array.

    Returns:
        The instantiated parameter.
    """    
    if is_variable(value):
        return value
    return Real(value)

parax.as_fixed(value)

Returns value as a parax.Fixed variable, wrapping it if necessary.

Parameters:

Name	Type	Description	Default
value	Param	An arbitrary variable or array-like object.	required

Returns:

Type	Description
Fixed	A fixed version of the variable.

Source code in parax/variables.py

def as_fixed(value: Param) -> Fixed:
    """
    Returns `value` as a `parax.Fixed` variable, wrapping it if necessary.

    Args:
        value: An arbitrary variable or array-like object.

    Returns:
        A fixed version of the variable.
    """    
    if isinstance(value, Fixed):
        return value
    return Fixed(value)

parax.variables.constrain_param(variable, *constraints)

Intelligently applies a constraint to a parameter (variable or array).

This function acts as a smart router for applying physical bounds to variables, regardless of how heavily wrapped they are. It safely drills through non-constrainable wrappers (like Fixed or Tagged), promotes unconstrained bases (like Real or raw JAX arrays), and correctly propagates constraints backwards through bijective transformations.

Parameters:

Name	Type	Description	Default
variable	Param	The target variable or standard JAX inexact array.	required
*constraints	AbstractConstraint	The physical constraints to apply.	()

Returns:

Name	Type	Description
Param	AbstractConstrainable	A new instance of the variable with the constraint applied.

Raises:

Type	Description
TypeError	If the variable type is not supported for dynamic constraining.

Source code in parax/variables.py

def constrain_param(variable: Param, *constraints: AbstractConstraint) -> AbstractConstrainable:
    """Intelligently applies a constraint to a parameter (variable or array).

    This function acts as a smart router for applying physical bounds to variables,
    regardless of how heavily wrapped they are. It safely drills through non-constrainable
    wrappers (like `Fixed` or `Tagged`), promotes unconstrained bases (like `Real` or
    raw JAX arrays), and correctly propagates constraints backwards through bijective
    transformations.

    Args:
        variable (Param): The target variable or standard JAX inexact array.
        *constraints (AbstractConstraint): The physical constraints to apply.

    Returns:
        Param: A new instance of the variable with the constraint applied.

    Raises:
        TypeError: If the variable type is not supported for dynamic constraining.
    """
    constraint = reduce(intersect_constraints, constraints)

    if is_constrainable(variable):
        return variable.constrain(constraint)
    elif isinstance(variable, Transformed):
        try:
            from distreqx.bijectors import Inverse
        except:
            from parax._bijectors import Inverse
        inverse_bij = Inverse(as_unwrapped(variable.bijector))
        transformed_constraint = TransformedConstraint(constraint, inverse_bij)
        new_inner = constrain_param(variable.raw_value, transformed_constraint)
        return Transformed(variable.bijector, new_inner)
    elif isinstance(variable, (Fixed, Tagged)):
        new_raw = constrain_param(variable.raw_value, constraint)
        return eqx.tree_at(lambda x: x.raw_value, variable, new_raw)
    elif isinstance(variable, Real) or eqx.is_inexact_array(variable):
        return Constrained(constraint, value=jnp.array(variable))
    else:
        raise TypeError(
            f"Cannot dynamically inject a constraint into type {type(variable)}. "
            "Ensure the target is a valid Parax variable or JAX inexact array."
        )