https://github.com/gerlero/parametrix

🪄 Flax-like computed parameters for bare JAX (and Equinox)
https://github.com/gerlero/parametrix

deep-learning equinox jax machine-learning neural-networks parameter parameter-estimation parameters physics-informed-neural-networks typed

Last synced: about 1 month ago
JSON representation

🪄 Flax-like computed parameters for bare JAX (and Equinox)

• Host: GitHub
• URL: https://github.com/gerlero/parametrix
• Owner: gerlero
• License: apache-2.0
• Created: 2025-05-21T00:13:28.000Z (about 1 year ago)
• Default Branch: main
• Last Pushed: 2026-04-28T18:51:37.000Z (2 months ago)
• Last Synced: 2026-04-28T20:33:29.431Z (2 months ago)
• Topics: deep-learning, equinox, jax, machine-learning, neural-networks, parameter, parameter-estimation, parameters, physics-informed-neural-networks, typed
• Language: Python
• Homepage: https://parametrix.readthedocs.io
• Size: 880 KB
• Stars: 0
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE.txt

Awesome Lists containing this project

README

          




**[`flax.nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param)-like computed parameters for bare [JAX](https://github.com/jax-ml/jax) (and [Equinox](https://github.com/patrick-kidger/equinox)).**

[![Documentation](https://img.shields.io/readthedocs/parametrix)](https://parametrix.readthedocs.io/)

[![CI](https://github.com/gerlero/parametrix/actions/workflows/ci.yml/badge.svg)](https://github.com/gerlero/parametrix/actions/workflows/ci.yml)

[![Codecov](https://codecov.io/gh/gerlero/parametrix/branch/main/graph/badge.svg)](https://codecov.io/gh/gerlero/parametrix)

[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

[![ty](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ty/main/assets/badge/v0.json)](https://github.com/astral-sh/ty)

[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)

[![Publish](https://github.com/gerlero/parametrix/actions/workflows/pypi-publish.yml/badge.svg)](https://github.com/gerlero/parametrix/actions/workflows/pypi-publish.yml)

[![PyPI](https://img.shields.io/pypi/v/parametrix)](https://pypi.org/project/parametrix/)

[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/parametrix)](https://pypi.org/project/parametrix/)



## Installation

```bash

pip install parametrix

```

## Example

The following example shows how to use [`Param`](https://parametrix.readthedocs.io) as a base class for a parameter class that enforces positivity:

```python

import jax.numpy as jnp

from parametrix import Param

class PositiveOnlyParam(Param):

    def __init__(self, value):

        super().__init__(jnp.log(value))

    @property

    def value(self):

        return jnp.exp(self.raw_value)

```

The backing values of `Param`s are always stored as `jax.Array`s, meaning that they will automatically be picked up as learnable parameters by libraries like Equinox.

`Param` objects also behave like numeric types, so that they are able to be used within models and any other functions without having to make any changes to the code.

## Documentation

API documentation is available at [Read the Docs](https://parametrix.readthedocs.io/).