parameter.Parameter

parameter.Parameter(value)

A minimal JAX-native parameter with full arithmetic support.

If placed at a position in the state tree, that position will take part in optimizations and will be differentiated. A state can be split into parameters and static parts by partition_state and combined by combine_state.

.value is the raw differentiable leaf updated by the optimizer. The value used in computation is .constrained_value (or collect_parameters); for plain Parameter the two coincide, but subclasses differ.

Parameters

Name	Type	Description	Default
value	Union[float, int, jnp.ndarray]	The parameter value. Will be converted to JAX array.	required

Examples

>>> p1 = Parameter(1.0)
>>> p2 = Parameter(jnp.array([1, 2, 3]))
>>> result = p1 + p2  # Works with JAX arithmetic
>>> grad_fn = jax.grad(lambda p: jnp.sum(p**2))
>>> gradients = grad_fn(p1)  # Seamless gradients

Attributes

Name	Description
constrained_value	Value as used in computation (post-transform / post-clip).
dtype	Data type of the parameter value.
ndim	Number of dimensions.
shape	Shape of the parameter value.
size	Total number of elements.

Methods

Name	Description
tree_flatten	Flatten for JAX pytree registration.
tree_unflatten	Unflatten for JAX pytree registration.

tree_flatten

parameter.Parameter.tree_flatten()

Flatten for JAX pytree registration.

tree_unflatten

parameter.Parameter.tree_unflatten(aux_data, children)

Unflatten for JAX pytree registration.

# parameter.Parameter { #tvboptim.parameter.Parameter } ```python parameter.Parameter(value) ``` A minimal JAX-native parameter with full arithmetic support. If placed at a position in the state tree, that position will take part in optimizations and will be differentiated. A state can be split into parameters and static parts by partition_state and combined by combine_state. ``.value`` is the raw differentiable leaf updated by the optimizer. The value used in computation is ``.constrained_value`` (or ``collect_parameters``); for plain ``Parameter`` the two coincide, but subclasses differ. ## Parameters {.doc-section .doc-section-parameters} | Name | Type | Description | Default | |--------|--------------------------------------------------------------------------------------------------------------|------------------------------------------------------|------------| | value | [Union](`typing.Union`)\[[float](`float`), [int](`int`), [jnp](`jax.numpy`).[ndarray](`jax.numpy.ndarray`)\] | The parameter value. Will be converted to JAX array. | _required_ | ## Examples {.doc-section .doc-section-examples} ```python >>> p1 = Parameter(1.0) >>> p2 = Parameter(jnp.array([1, 2, 3])) >>> result = p1 + p2 # Works with JAX arithmetic >>> grad_fn = jax.grad(lambda p: jnp.sum(p**2)) >>> gradients = grad_fn(p1) # Seamless gradients ``` ## Attributes | Name | Description | | --- | --- | | [constrained_value](#tvboptim.parameter.Parameter.constrained_value) | Value as used in computation (post-transform / post-clip). | | [dtype](#tvboptim.parameter.Parameter.dtype) | Data type of the parameter value. | | [ndim](#tvboptim.parameter.Parameter.ndim) | Number of dimensions. | | [shape](#tvboptim.parameter.Parameter.shape) | Shape of the parameter value. | | [size](#tvboptim.parameter.Parameter.size) | Total number of elements. | ## Methods | Name | Description | | --- | --- | | [tree_flatten](#tvboptim.parameter.Parameter.tree_flatten) | Flatten for JAX pytree registration. | | [tree_unflatten](#tvboptim.parameter.Parameter.tree_unflatten) | Unflatten for JAX pytree registration. | ### tree_flatten { #tvboptim.parameter.Parameter.tree_flatten } ```python parameter.Parameter.tree_flatten() ``` Flatten for JAX pytree registration. ### tree_unflatten { #tvboptim.parameter.Parameter.tree_unflatten } ```python parameter.Parameter.tree_unflatten(aux_data, children) ``` Unflatten for JAX pytree registration.