types.stateutils

types.stateutils

Functions

Name	Description
collect_parameters	Extract values and modules from parameter owners in a state tree.
combine_state	Recombine optimized parameters with static values.
is_parameter_owner	Whether value explicitly owns trainable parameter leaves.
is_partition_leaf	Whether state partition/combine must treat value atomically.
mark_parameters	Mark ordinary parameters and nested Equinox array leaves.
parameter_filter	Build the nested differentiability filter for one parameter owner.
partition_state	Separate explicitly trainable leaves from static state.
show_parameters	Show explicitly trainable parameter leaves in the tree.

collect_parameters

types.stateutils.collect_parameters(state)

Extract values and modules from parameter owners in a state tree.

This function traverses a JAX PyTree state and extracts the underlying values from Parameter objects and modules from EquinoxParameter objects while leaving other values unchanged.

Parameters

Name	Type	Description	Default
state	Any	JAX PyTree containing Parameter objects and other values.	required

Returns

Name	Type	Description
	Any	JAX PyTree with same structure as input, but Parameter objects replaced by their underlying values.

Examples

>>> from tvboptim.types import Parameter
>>> import jax.numpy as jnp
>>>
>>> # Create state with Parameter objects
>>> state = {
...     'param1': Parameter(jnp.array(1.5)),
...     'param2': jnp.array(2.0),
...     'nested': {'param3': Parameter(jnp.array([1, 2, 3]))}
... }
>>>
>>> # Extract values
>>> values = collect_parameters(state)
>>> print(values['param1'])  # jnp.array(1.5)
>>> print(values['param2'])  # jnp.array(2.0)

Notes

This function is useful when you need to extract raw JAX arrays from a state tree for operations that don’t require Parameter metadata. With the new Parameter system, this function may become less necessary as Parameters support the JAX array protocol directly.

combine_state

types.stateutils.combine_state(diff_state, static_state)

Recombine optimized parameters with static values.

is_parameter_owner

types.stateutils.is_parameter_owner(value)

Whether value explicitly owns trainable parameter leaves.

is_partition_leaf

types.stateutils.is_partition_leaf(value)

Whether state partition/combine must treat value atomically.

mark_parameters

types.stateutils.mark_parameters(state)

Mark ordinary parameters and nested Equinox array leaves.

parameter_filter

types.stateutils.parameter_filter(value)

Build the nested differentiability filter for one parameter owner.

partition_state

types.stateutils.partition_state(state)

Separate explicitly trainable leaves from static state.

show_parameters

types.stateutils.show_parameters(tree)

Show explicitly trainable parameter leaves in the tree.