### Install Paramax

Source: https://github.com/danielward27/paramax/blob/main/docs/index.md

Install the Paramax package using pip.

```bash
pip install paramax
```

--------------------------------

### Run Documentation Doctests

Source: https://github.com/danielward27/paramax/blob/main/docs/README.md

Executes doctest code blocks to ensure documentation examples are functional.

```bash
make -C docs doctest
```

--------------------------------

### Build Documentation with Sphinx

Source: https://github.com/danielward27/paramax/blob/main/docs/README.md

Generates HTML documentation from the docs directory.

```bash
make -C docs html
```

--------------------------------

### Initialize positive parameter to target value

Source: https://context7.com/danielward27/paramax/llms.txt

Use inv_softplus to map a target scale value to a parameter space that results in the desired scale after a softplus transformation.

```python
target_scale = jnp.array([0.5, 1.0, 2.0])
init_param = inv_softplus(target_scale)
# softplus(init_param) ≈ target_scale
```

--------------------------------

### Create Custom Parameterization Wrapper

Source: https://context7.com/danielward27/paramax/llms.txt

Base class for creating custom parameterization wrappers. Implement the `unwrap()` method to define custom constraint behavior.

```python
import paramax
import jax.numpy as jnp
import equinox as eqx
from jaxtyping import Array

# Custom wrapper for symmetric positive definite matrices
class SymmetricPositiveDefinite(paramax.AbstractUnwrappable[Array]):
    """Parameterize a symmetric positive definite matrix via Cholesky factor."""
    lower_triangular: Array

    def __init__(self, size: int, key):
        import jax.random as jr
        # Initialize with small random values
        n_params = size * (size + 1) // 2
        params = jr.normal(key, (n_params,)) * 0.1
        idx = jnp.tril_indices(size)
        self.lower_triangular = jnp.zeros((size, size)).at[idx].set(params)

    def unwrap(self) -> Array:
        # L @ L.T is guaranteed symmetric positive definite
        L = self.lower_triangular
        return L @ L.T

# Usage
import jax.random as jr
spd = SymmetricPositiveDefinite(3, jr.key(0))
matrix = paramax.unwrap(spd)
# matrix is 3x3 symmetric positive definite

# Verify positive definiteness
eigenvalues = jnp.linalg.eigvalsh(matrix)
# All eigenvalues > 0
```

--------------------------------

### Parameterize Callable with Trainable Parameters

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Unwraps an object by calling a function with provided arguments and keyword arguments. All components (function, args, kwargs) can contain trainable parameters. Recommended for vectorized contexts to ensure correct unwrapping.

```python
class Parameterize(fn: Callable, *args, **kwargs)
```

--------------------------------

### AbstractUnwrappable

Source: https://context7.com/danielward27/paramax/llms.txt

Base class for creating custom parameterization wrappers. Subclass it and implement the `unwrap()` method.

```APIDOC
## AbstractUnwrappable

### Description
The `AbstractUnwrappable` base class serves as a foundation for creating custom parameterization wrappers within Paramax. By subclassing `AbstractUnwrappable` and implementing the `unwrap()` method, users can define custom constraint behaviors for their parameters.

### Method
Inheritance and implementation of the `unwrap()` method.

### Endpoint
N/A (Class definition)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None (This is a base class for custom implementations)

### Request Example
```python
import paramax
import jax.numpy as jnp
import equinox as eqx
from jaxtyping import Array
import jax.random as jr

# Custom wrapper for symmetric positive definite matrices
class SymmetricPositiveDefinite(paramax.AbstractUnwrappable[Array]):
    """Parameterize a symmetric positive definite matrix via Cholesky factor."""
    lower_triangular: Array

    def __init__(self, size: int, key):
        # Initialize with small random values for the lower triangular part
        n_params = size * (size + 1) // 2
        params = jr.normal(key, (n_params,)) * 0.1
        idx = jnp.tril_indices(size)
        self.lower_triangular = jnp.zeros((size, size)).at[idx].set(params)

    def unwrap(self) -> Array:
        # L @ L.T is guaranteed symmetric positive definite
        L = self.lower_triangular
        return L @ L.T

# Usage
key = jr.key(0)
spd = SymmetricPositiveDefinite(3, key)
matrix = paramax.unwrap(spd)
# matrix is a 3x3 symmetric positive definite matrix

# Verify positive definiteness
eigenvalues = jnp.linalg.eigvalsh(matrix)
# All eigenvalues should be > 0
```

### Response
#### Success Response (200)
Returns the unwrapped, constrained parameter value.

#### Response Example
```python
# The 'matrix' variable in the example above holds the unwrapped SPD matrix.
print(matrix)
```
```

--------------------------------

### Parameterize class

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Unwraps an object by calling a function with provided arguments and keyword arguments. The function, arguments, and keyword arguments can all contain trainable parameters.

```APIDOC
## class Parameterize(fn: Callable, *args, **kwargs)

### Description
Unwrap an object by calling fn with args and kwargs. All of fn, args and kwargs may contain trainable parameters.

#### NOTE
Unwrapping typically occurs after model initialization. Therefore, if the `Parameterize` object may be created in a vectorized context, we recommend ensuring that `fn` still unwraps correctly, e.g. by supporting broadcasting.

### Parameters
* **fn** – Callable to call with args, and kwargs.
* **\*args** – Positional arguments to pass to fn.
* **\*\*kwargs** – Keyword arguments to pass to fn.

#### fn *: Callable[[...], T]*

#### args *: tuple[Any, ...]*

#### kwargs *: dict[str, Any]*

### Methods
#### unwrap() -> T
```

--------------------------------

### Apply Weight Normalization

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Applies weight normalization to a weight matrix, which may be wrapped. This technique is based on the paper https://arxiv.org/abs/1602.07868.

```python
class WeightNormalization(weight: Array | [AbstractUnwrappable](#paramax.wrappers.AbstractUnwrappable)[Array])
```

--------------------------------

### Apply Positivity Constraint

Source: https://github.com/danielward27/paramax/blob/main/README.md

Use Parameterize to enforce positivity on a PyTree component and unwrap it for use.

```python
>>> import paramax
>>> import jax.numpy as jnp
>>> scale = paramax.Parameterize(jnp.exp, jnp.log(jnp.ones(3)))  # Enforce positivity
>>> paramax.unwrap(("abc", 1, scale))
('abc', 1, Array([1., 1., 1.], dtype=float32))
```

--------------------------------

### Parameterize Class

Source: https://context7.com/danielward27/paramax/llms.txt

The `Parameterize` class is a flexible wrapper that calls a specified function with stored arguments and keyword arguments when unwrapped. The function, args, and kwargs can all contain trainable parameters.

```APIDOC
## Parameterize Class

### Description
A flexible wrapper that calls a specified function with stored arguments and keyword arguments when unwrapped. The function, args, and kwargs can all contain trainable parameters.

### Usage
```python
import paramax
import jax.numpy as jnp
from jax.nn import softplus

# Enforce positivity with exponential
positive_params = paramax.Parameterize(jnp.exp, jnp.zeros(3))
paramax.unwrap(positive_params)
# Array([1., 1., 1.], dtype=float32)

# Create diagonal matrix from vector
diag_matrix = paramax.Parameterize(jnp.diag, jnp.ones(3))
paramax.unwrap(diag_matrix)
# Array([[1., 0., 0.],
#        [0., 1., 0.],
#        [0., 0., 1.]], dtype=float32)

# Use softplus for smooth positivity constraint
smooth_positive = paramax.Parameterize(softplus, jnp.array([-1.0, 0.0, 1.0]))
paramax.unwrap(smooth_positive)
# Array with positive values transformed via softplus

# Create lower triangular matrix
def lower_triangular(params, size):
    idx = jnp.tril_indices(size)
    return jnp.zeros((size, size)).at[idx].set(params)

n = 3
n_params = n * (n + 1) // 2  # 6 parameters for 3x3 lower triangular
lower_tri = paramax.Parameterize(lower_triangular, jnp.arange(n_params, dtype=float), size=n)
paramax.unwrap(lower_tri)
# Array([[0., 0., 0.],
#        [1., 2., 0.],
#        [3., 4., 5.]], dtype=float32)
```
```

--------------------------------

### RealToIncreasingOnInterval class

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Maps an unconstrained vector to increasing points on a fixed interval. The input vector is transformed via a softmax into positive widths, which fill the interval after adding a minimum width.

```APIDOC
## class RealToIncreasingOnInterval(arr: Array, interval: tuple[float | int, float | int], min_width: float, include_endpoints: Literal['both', 'neither', 'lower', 'upper'])

### Description
Map an unconstrained vector to increasing points on a fixed interval. The input vector is transformed via a softmax into positive widths, to fill the interval after adding minimum width. The cumulative sum of the widths produces the increasing points. If an array of size d is used, the result has size d+1 if both endpoints are included, d if one is included, and d-1 if neither are included.

### Parameters
* **arr** – Unconstrained vector parameterizing the widths.
* **interval** – (lower, upper) bounds of the interval.
* **min_width** – Minimum spacing between consecutive points.
* **include_endpoints** – Which endpoints to include: “both”, “neither”, “lower”, or “upper”.

#### arr *: Array*

#### interval *: tuple[float | int, float | int]*

#### min_width *: float*

#### include_endpoints *: Literal['both', 'neither', 'lower', 'upper']* 

### Methods
#### unwrap() -> Array
```

--------------------------------

### Inverse Softplus Utility Function

Source: https://context7.com/danielward27/paramax/llms.txt

Computes the inverse of the softplus function. Useful for initializing parameters that will be transformed via softplus to ensure positivity.

```python
import jax.numpy as jnp
from jax.nn import softplus
from paramax.utils import inv_softplus

# Compute inverse softplus
x = jnp.array([1.0, 2.0, 3.0])
y = softplus(x)
x_reconstructed = inv_softplus(y)
# x_reconstructed ≈ x (within numerical precision)
```

--------------------------------

### Freeze Parameters with NonTrainable Wrapper

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Wraps inexact array leaves with NonTrainable to freeze parameters. This is generally preferable to wrapping the entire tree, allowing easier attribute access. Note that underlying parameters might still be affected by regularization.

```python
non_trainable(tree: PyTree)
```

```python
class NonTrainable(tree: T)
```

```python
>>> eqx.partition(
...     ...,
...     is_leaf=lambda leaf: isinstance(leaf, wrappers.NonTrainable),
... )
```

--------------------------------

### Map Vector to Increasing Points on Interval

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Transforms an unconstrained vector into increasing points on a fixed interval using a softmax-based width calculation. The output size depends on whether endpoints are included.

```python
class RealToIncreasingOnInterval(arr: Array, interval: tuple[float | int, float | int], , min_width: float, include_endpoints: Literal['both', 'neither', 'lower', 'upper'])
```

--------------------------------

### inv_softplus

Source: https://github.com/danielward27/paramax/blob/main/docs/api/utils.md

Calculates the inverse of the softplus function for a given input array.

```APIDOC
## inv_softplus

### Description
The inverse of the softplus function, checking for positive inputs.

### Parameters
- **x** (ArrayLike) - Required - The input array to process.

### Response
- **Returns** (Array) - The result of the inverse softplus operation.
```

--------------------------------

### Apply Weight Normalization

Source: https://context7.com/danielward27/paramax/llms.txt

Applies weight normalization to a weight matrix. The learned scale parameter approximates the row norms of the normalized weights. Can be nested with other wrappers.

```python
import jax.numpy as jnp
import jax.random as jr
import paramax

key = jr.key(0)
weight_matrix = jr.normal(key, (10, 3))
weight_norm = paramax.WeightNormalization(weight_matrix)

# Unwrap to get normalized weights
normalized = paramax.unwrap(weight_norm)
# Shape: (10, 3) - same as input

# The row norms equal the learned scale parameter
scales = paramax.unwrap(weight_norm.scale)
row_norms = jnp.linalg.norm(normalized, axis=-1, keepdims=True)
# scales ≈ row_norms (within numerical precision)

# Can be nested with other wrappers
positive_weights = paramax.Parameterize(jnp.exp, jr.normal(key, (5, 4)))
normalized_positive = paramax.WeightNormalization(positive_weights)
result = paramax.unwrap(normalized_positive)
# Weights are first made positive, then normalized
```

--------------------------------

### Map Real to Increasing Points on Interval

Source: https://context7.com/danielward27/paramax/llms.txt

Maps unconstrained vectors to monotonically increasing points within a specified interval. Supports including/excluding endpoints and non-uniform spacing.

```python
import paramax
import jax.numpy as jnp

# Map 2 parameters to 3 increasing points in [-1, 2], including both endpoints
increasing = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="both"
)
paramax.unwrap(increasing)
# Array([-1. ,  0.5,  2. ], dtype=float32)
```

```python
# Include only lower endpoint
lower_only = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="lower"
)
paramax.unwrap(lower_only)
# Array([-1. ,  0.5], dtype=float32)
```

```python
# Include only upper endpoint
upper_only = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="upper"
)
paramax.unwrap(upper_only)
# Array([0.5, 2. ], dtype=float32)
```

```python
# Neither endpoint included
neither = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="neither"
)
paramax.unwrap(neither)
# Array([0.5], dtype=float32)
```

```python
# Non-uniform spacing via different input values
skewed = paramax.RealToIncreasingOnInterval(
    jnp.array([-100.0, 0.0]),  # Strong bias toward first width being minimal
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="both"
)
paramax.unwrap(skewed)
# Array([-1. , -0.9,  2. ], dtype=float32) - points clustered near lower bound
```

--------------------------------

### non_trainable function

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Freezes parameters by wrapping inexact array leaves with NonTrainable. This is generally preferable to wrapping the entire tree.

```APIDOC
### non_trainable(tree: PyTree)

### Description
Freezes parameters by wrapping inexact array leaves with [`NonTrainable`](#paramax.wrappers.NonTrainable). Regularization is likely to apply before unwrapping. To avoid regularization impacting non-trainable parameters, they should be filtered out, for example using `eqx.partition`.

#### NOTE
Wrapping the arrays in a model rather than the entire tree is often preferable, allowing easier access to attributes compared to wrapping the entire tree.

### Parameters
* **tree** – The pytree.

### Example
```python
>>> import equinox as eqx
>>> import jax
>>> import jax.numpy as jnp
>>> import paramax.wrappers as wrappers

>>> x = jnp.array([1.0, 2.0])
>>> wrapped_x = wrappers.non_trainable(x)
>>> print(wrapped_x)
NonTrainable(tree=DeviceArray([1., 2.], dtype=float32))

>>> # Example of filtering during partitioning
>>> model = {'a': jnp.array(1.0), 'b': wrappers.NonTrainable(jnp.array(2.0))}
>>> treedef, leaves = jax.tree_util.tree_flatten(model)
>>> partitioned_model = eqx.partition(
...     model,
...     is_leaf=lambda leaf: isinstance(leaf, wrappers.NonTrainable)
... )
>>> print(partitioned_model)
(PyTreeDef(dict<leaf=None, a=None, b=None>),
 PyTreeDef(dict<leaf=None, a=None, b=None>))
```
```

--------------------------------

### Parameterize PyTree Nodes

Source: https://context7.com/danielward27/paramax/llms.txt

The Parameterize class applies a function to stored arguments upon unwrapping, useful for enforcing constraints like positivity or creating structured matrices.

```python
import paramax
import jax.numpy as jnp
from jax.nn import softplus

# Enforce positivity with exponential
positive_params = paramax.Parameterize(jnp.exp, jnp.zeros(3))
paramax.unwrap(positive_params)
# Array([1., 1., 1.], dtype=float32)

# Create diagonal matrix from vector
diag_matrix = paramax.Parameterize(jnp.diag, jnp.ones(3))
paramax.unwrap(diag_matrix)
# Array([[1., 0., 0.],
#        [0., 1., 0.],
#        [0., 0., 1.]], dtype=float32)

# Use softplus for smooth positivity constraint
smooth_positive = paramax.Parameterize(softplus, jnp.array([-1.0, 0.0, 1.0]))
paramax.unwrap(smooth_positive)
# Array with positive values transformed via softplus

# Create lower triangular matrix
def lower_triangular(params, size):
    idx = jnp.tril_indices(size)
    return jnp.zeros((size, size)).at[idx].set(params)

n = 3
n_params = n * (n + 1) // 2  # 6 parameters for 3x3 lower triangular
lower_tri = paramax.Parameterize(lower_triangular, jnp.arange(n_params, dtype=float), size=n)
paramax.unwrap(lower_tri)
# Array([[0., 0., 0.],
#        [1., 2., 0.],
#        [3., 4., 5.]], dtype=float32)
```

--------------------------------

### Freeze Parameters with NonTrainable

Source: https://context7.com/danielward27/paramax/llms.txt

The NonTrainable class applies stop_gradient to array leaves, preventing them from being updated during gradient computation.

```python
import paramax
import jax.numpy as jnp
import equinox as eqx

# Wrap a single array as non-trainable
frozen_weights = paramax.NonTrainable(jnp.ones(3))
paramax.unwrap(frozen_weights)
# Array([1., 1., 1.], dtype=float32) - but gradients will be zero

# Use in a loss function to verify gradients are blocked
model = (jnp.ones(3), jnp.array(2.0))
frozen_model = (paramax.NonTrainable(model[0]), model[1])

def loss(m):
    m = paramax.unwrap(m)
    return m[0].sum() + m[1]

grads = eqx.filter_grad(loss)(frozen_model)
# grads[0] will be zero (frozen), grads[1] will be 1.0 (trainable)
```

--------------------------------

### RealToIncreasingOnInterval

Source: https://context7.com/danielward27/paramax/llms.txt

Maps an unconstrained vector to monotonically increasing points within a specified interval, useful for parameterizing ordered sequences.

```APIDOC
## RealToIncreasingOnInterval

### Description
The `RealToIncreasingOnInterval` class maps an unconstrained vector to monotonically increasing points within a specified interval. This is useful for parameterizing knot positions, quantiles, or any ordered sequence where the values must strictly increase.

### Method
Instantiation of the `RealToIncreasingOnInterval` class.

### Endpoint
N/A (Class definition)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **unconstrained_vector** (Array) - The input unconstrained vector.
- **interval** (tuple) - A tuple `(min_val, max_val)` defining the interval.
- **min_width** (float, optional) - The minimum width between consecutive points. Defaults to 0.1.
- **include_endpoints** (str, optional) - Specifies whether to include interval endpoints. Options: "both", "lower", "upper", "neither". Defaults to "both".

### Request Example
```python
import paramax
import jax.numpy as jnp

# Map 2 parameters to 3 increasing points in [-1, 2], including both endpoints
increasing = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="both"
)
paramax.unwrap(increasing)
# Array([-1. ,  0.5,  2. ], dtype=float32)
```

### Response
#### Success Response (200)
Returns an array of monotonically increasing points within the specified interval.

#### Response Example
```python
# Example for include_endpoints="lower"
lower_only = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="lower"
)
paramax.unwrap(lower_only)
# Array([-1. ,  0.5], dtype=float32)

# Example for include_endpoints="upper"
upper_only = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="upper"
)
paramax.unwrap(upper_only)
# Array([0.5, 2. ], dtype=float32)

# Example for include_endpoints="neither"
neither = paramax.RealToIncreasingOnInterval(
    jnp.zeros(2),
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="neither"
)
paramax.unwrap(neither)
# Array([0.5], dtype=float32)

# Example with skewed input values
skewed = paramax.RealToIncreasingOnInterval(
    jnp.array([-100.0, 0.0]),  # Strong bias toward first width being minimal
    interval=(-1, 2),
    min_width=0.1,
    include_endpoints="both"
)
paramax.unwrap(skewed)
# Array([-1. , -0.9,  2. ], dtype=float32) - points clustered near lower bound
```
```

--------------------------------

### WeightNormalization class

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Applies weight normalization to a weight matrix, as described in https://arxiv.org/abs/1602.07868.

```APIDOC
## class WeightNormalization(weight: Array | [AbstractUnwrappable](#paramax.wrappers.AbstractUnwrappable)[Array])

### Description
Applies weight normalization ([https://arxiv.org/abs/1602.07868](https://arxiv.org/abs/1602.07868)).

### Parameters
* **weight** – The (possibly wrapped) weight matrix.

#### weight *: Array | [AbstractUnwrappable](#paramax.wrappers.AbstractUnwrappable)[Array]*

### Attributes
#### scale *: Array | [AbstractUnwrappable](#paramax.wrappers.AbstractUnwrappable)[Array]*

### Methods
#### unwrap() -> Array
```

--------------------------------

### NonTrainable Class

Source: https://context7.com/danielward27/paramax/llms.txt

The `NonTrainable` class wraps a PyTree and applies `stop_gradient` to all array-like leaves when unwrapped, effectively freezing parameters during gradient computation.

```APIDOC
## NonTrainable Class

### Description
Wraps a PyTree and applies `stop_gradient` to all array-like leaves when unwrapped, effectively freezing parameters during gradient computation.

### Usage
```python
import paramax
import jax.numpy as jnp
import equinox as eqx

# Wrap a single array as non-trainable
frozen_weights = paramax.NonTrainable(jnp.ones(3))
paramax.unwrap(frozen_weights)
# Array([1., 1., 1.], dtype=float32) - but gradients will be zero

# Use in a loss function to verify gradients are blocked
model = (jnp.ones(3), jnp.array(2.0))
frozen_model = (paramax.NonTrainable(model[0]), model[1])

def loss(m):
    m = paramax.unwrap(m)
    return m[0].sum() + m[1]

grads = eqx.filter_grad(loss)(frozen_model)
# grads[0] will be zero (frozen), grads[1] will be 1.0 (trainable)
```
```

--------------------------------

### WeightNormalization

Source: https://context7.com/danielward27/paramax/llms.txt

Applies weight normalization to a weight matrix, learning a scale parameter for each row. Can be nested with other wrappers.

```APIDOC
## WeightNormalization

### Description
Applies weight normalization to a weight matrix. It learns a scale parameter for each row, ensuring that the row norms are close to these learned scales.

### Method
Instantiation of the `WeightNormalization` class.

### Endpoint
N/A (Class definition)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **weight_matrix** (Array) - The input weight matrix to be normalized.

### Request Example
```python
import paramax
import jax.numpy as jnp
import jax.random as jr

key = jr.key(0)
weight_matrix = jr.normal(key, (10, 3))
weight_norm = paramax.WeightNormalization(weight_matrix)
```

### Response
#### Success Response (200)
Returns a `WeightNormalization` object.

#### Response Example
```json
{
  "normalized_weights": "(10, 3) array",
  "scale": "(10, 1) array"
}
```

### Nested Usage Example
```python
import paramax
import jax.numpy as jnp
import jax.random as jr

key = jr.key(0)
positive_weights = paramax.Parameterize(jnp.exp, jr.normal(key, (5, 4)))
normalized_positive = paramax.WeightNormalization(positive_weights)
result = paramax.unwrap(normalized_positive)
# Weights are first made positive, then normalized
```
```

--------------------------------

### non_trainable() Function

Source: https://context7.com/danielward27/paramax/llms.txt

The `non_trainable()` function wraps all inexact array leaves in a PyTree with `NonTrainable`, providing a convenient way to freeze an entire subtree or model.

```APIDOC
## non_trainable() Function

### Description
Wraps all inexact array leaves in a PyTree with `NonTrainable`, providing a convenient way to freeze an entire subtree or model.

### Usage
```python
import paramax
import jax.numpy as jnp
import equinox as eqx

# Freeze all array parameters in a tuple
model = (jnp.ones(3), jnp.array(2.0), "metadata")
frozen = paramax.non_trainable(model)
# Only arrays are wrapped: (NonTrainable([1,1,1]), NonTrainable(2.0), "metadata")

# Verify gradients are blocked
def loss(m):
    m = paramax.unwrap(m)
    return m[0].sum() + m[1]

grads = eqx.filter_grad(loss)(frozen)
# All array gradients will be zero

# Filter non-trainable params for regularization exclusion
trainable, non_trainable_params = eqx.partition(
    frozen,
    is_leaf=lambda leaf: isinstance(leaf, paramax.NonTrainable),
    filter_spec=lambda leaf: not isinstance(leaf, paramax.NonTrainable),
)
```
```

--------------------------------

### Weight Normalization

Source: https://context7.com/danielward27/paramax/llms.txt

WeightNormalization implements reparameterization to decouple magnitude from direction.

```python
import paramax
import jax.numpy as jnp
import jax.random as jr
```

--------------------------------

### NonTrainable class

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Applies stop gradient to all arraylike leaves before unwrapping. Useful for marking pytrees as frozen/non-trainable, though non_trainable() is often preferred.

```APIDOC
## class NonTrainable(tree: T)

### Description
Applies stop gradient to all arraylike leaves before unwrapping. See also [`non_trainable()`](#paramax.wrappers.non_trainable), which is probably a generally prefereable way to achieve similar behaviour, which wraps the arraylike leaves directly, rather than the tree. Useful to mark pytrees (arrays, submodules, etc) as frozen/non-trainable. Note that the underlying parameters may still be impacted by regularization, so it is generally advised to use this as a suggestively named class for filtering parameters.

### Parameters
* **tree** – The pytree.

#### tree *: T*

### Methods
#### unwrap() -> T
```

--------------------------------

### AbstractUnwrappable

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

An abstract class representing an unwrappable object. Unwrappables replace PyTree nodes, applying custom behavior upon unwrapping.

```APIDOC
## class AbstractUnwrappable

### Description
An abstract class representing an unwrappable object. Unwrappables replace PyTree nodes, applying custom behavior upon unwrapping.

### Methods
#### abstractmethod unwrap() -> T
Returns the unwrapped pytree, assuming no wrapped subnodes exist.
```

--------------------------------

### Unwrap PyTree Nodes

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Maps across a PyTree and unwraps all AbstractUnwrappable nodes, leaving others unchanged. Innermost unwrappable nodes are unwrapped first.

```python
unwrap(tree: PyTree)
```

--------------------------------

### Freeze Subtrees with non_trainable

Source: https://context7.com/danielward27/paramax/llms.txt

The non_trainable function automatically wraps all inexact array leaves in a PyTree to freeze them.

```python
import paramax
import jax.numpy as jnp
import equinox as eqx

# Freeze all array parameters in a tuple
model = (jnp.ones(3), jnp.array(2.0), "metadata")
frozen = paramax.non_trainable(model)
# Only arrays are wrapped: (NonTrainable([1,1,1]), NonTrainable(2.0), "metadata")

# Verify gradients are blocked
def loss(m):
    m = paramax.unwrap(m)
    return m[0].sum() + m[1]

grads = eqx.filter_grad(loss)(frozen)
# All array gradients will be zero

# Filter non-trainable params for regularization exclusion
trainable, non_trainable_params = eqx.partition(
    frozen,
    is_leaf=lambda leaf: isinstance(leaf, paramax.NonTrainable),
    filter_spec=lambda leaf: not isinstance(leaf, paramax.NonTrainable),
)
```

--------------------------------

### unwrap function

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Maps across a PyTree and unwraps all AbstractUnwrappable nodes, leaving other nodes unchanged. Innermost AbstractUnwrappable nodes are unwrapped first.

```APIDOC
### unwrap(tree: PyTree)

### Description
Map across a PyTree and unwrap all [`AbstractUnwrappable`](#paramax.wrappers.AbstractUnwrappable) nodes. This leaves all other nodes unchanged. If nested, the innermost `AbstractUnwrappable` nodes are unwrapped first.

### Example
Enforcing positivity.
```

--------------------------------

### Unwrap PyTrees

Source: https://context7.com/danielward27/paramax/llms.txt

The unwrap function resolves AbstractUnwrappable placeholders within a PyTree, processing nested structures from the inside out.

```python
import paramax
import jax.numpy as jnp

# Enforce positivity using exponential transformation
scale = paramax.Parameterize(jnp.exp, jnp.log(jnp.ones(3)))
result = paramax.unwrap(("abc", 1, scale))
# Result: ('abc', 1, Array([1., 1., 1.], dtype=float32))

# Nested unwrapping - innermost first
nested = paramax.Parameterize(
    jnp.square,
    paramax.Parameterize(jnp.square, paramax.Parameterize(jnp.square, 2))
)
result = paramax.unwrap(nested)
# Result: 256 (equivalent to square(square(square(2))))
```

--------------------------------

### Check for Unwrappable Objects in PyTree

Source: https://context7.com/danielward27/paramax/llms.txt

Determines if a PyTree contains any `AbstractUnwrappable` objects. Useful for validation or conditional unwrapping logic.

```python
import paramax
import jax.numpy as jnp

# Check PyTree with unwrappable
model_with_params = {
    "weights": paramax.Parameterize(jnp.exp, jnp.zeros(3)),
    "bias": jnp.ones(3),
}
paramax.contains_unwrappables(model_with_params)
# True
```

```python
# Check PyTree without unwrappable
plain_model = {
    "weights": jnp.zeros(3),
    "bias": jnp.ones(3),
}
paramax.contains_unwrappables(plain_model)
# False
```

```python
# Useful for assertions or conditional unwrapping
def forward(model, x):
    if paramax.contains_unwrappables(model):
        model = paramax.unwrap(model)
    return model["weights"] @ x + model["bias"]
```

--------------------------------

### Check for Unwrappables in PyTree

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Determines if a given pytree contains any unwrappable nodes.

```python
contains_unwrappables(pytree)
```

--------------------------------

### contains_unwrappables function

Source: https://github.com/danielward27/paramax/blob/main/docs/api/wrappers.md

Checks if a given pytree contains any unwrappable nodes.

```APIDOC
### contains_unwrappables(pytree)

### Description
Check if a pytree contains unwrappable nodes.
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.