drinx 1.0.0__tar.gz

drinx-1.0.0/.github/actions/setup-python-env/action.yml

@@ -0,0 +1,30 @@
+name: "Setup Python Environment"
+description: "Set up Python environment for the given Python version"
+
+inputs:
+  python-version:
+    description: "Python version to use"
+    required: true
+    default: "3.12"
+  uv-version:
+    description: "uv version to use"
+    required: true
+    default: "0.5.8"
+
+runs:
+  using: "composite"
+  steps:
+    - uses: actions/setup-python@v5
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v2
+      with:
+        version: ${{ inputs.uv-version }}
+        enable-cache: "true"
+        cache-suffix: ${{ matrix.python-version }}
+
+    - name: Install Python dependencies
+      run: uv sync --extra dev
+      shell: bash

drinx-1.0.0/.github/workflows/cicd.yml

@@ -0,0 +1,72 @@
+name: CI-CD
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  checks:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - uses: actions/cache@v4
+        with:
+          path: ~/.cache/pre-commit
+          key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }}
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+
+      - name: Run pre-commit
+        run: uv run pre-commit run -a --show-diff-on-failure
+
+      - name: build docs
+        run: uv run sphinx-build -W --keep-going docs/source/ docs/build/
+
+  tests:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      id-token: write
+    strategy:
+      matrix:
+        python-version: ["3.11"]
+      fail-fast: false
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Run tests
+        run: uv run python -m pytest tests --cov --cov-branch --cov-config=pyproject.toml --cov-report=xml
+
+      - name: Typecheck with ty
+        run: uvx ty check --error-on-warning
+
+      - name: Upload coverage reports to Codecov
+        uses: codecov/codecov-action@v4
+        if: ${{ matrix.python-version == '3.11' }}
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: ./coverage.xml
+          fail_ci_if_error: true
+          slug: ymahlau/drinx

drinx-1.0.0/.github/workflows/publish.yml

@@ -0,0 +1,42 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+  workflow_dispatch:  # Allows manual triggering from Actions tab
+
+permissions:
+  contents: read
+
+jobs:
+  pypi-publish:
+    name: Build and publish Python 🐍 distributions to PyPI
+    runs-on: ubuntu-latest
+    # This permission is REQUIRED for Trusted Publishing
+    permissions:
+      id-token: write  
+      contents: read
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install build tools
+        run: pip install build
+
+      - name: Build package
+        run: python -m build
+
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          # We don't need a password/token input because we use 
+          # Trusted Publishing (id-token permission above).
+          print-hash: true

drinx-1.0.0/.gitignore

@@ -0,0 +1,9 @@
+uv.lock
+.venv
+outputs/
+*.pyc
+*.npz
+docs/build
+docs/jupyter_execute
+docs/source/api
+scripts_dev/

drinx-1.0.0/.pre-commit-config.yaml

@@ -0,0 +1,34 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-yaml
+      - id: check-added-large-files
+      - id: check-shebang-scripts-are-executable
+      - id: check-toml
+      - id: check-merge-conflict
+        args: [--assume-in-merge]
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.11.2
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args:
+          - --fix
+      # Run the formatter.
+      - id: ruff-format
+  - repo: https://github.com/woodruffw/zizmor-pre-commit
+    rev: v1.5.2
+    hooks:
+      - id: zizmor
+  - repo: local
+    hooks:
+      - id: ty
+        name: ty (via uvx)
+        entry: uvx ty check --error-on-warning
+        language: system
+        types: [python]
+        pass_filenames: false

drinx-1.0.0/.readthedocs.yml

@@ -0,0 +1,14 @@
+version: 2
+formats: [htmlzip]
+
+sphinx:
+  configuration: docs/source/conf.py
+
+build:
+  os: "ubuntu-22.04"
+  tools:
+    python: "3.12"
+  jobs:
+    install:
+      - pip install -r docs/requirements.txt
+      - pip install --no-deps .

drinx-1.0.0/CLAUDE.md

@@ -0,0 +1,61 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+Drinx is a small Python library ("Dataclass Registry in JAX") that wraps Python's standard `dataclasses` module to make dataclasses compatible with JAX transformations (e.g., `jit`, `vmap`, `grad`). It does this by registering each decorated class as a JAX pytree node, with support for marking fields as "static" (excluded from JAX tracing) or "dynamic" (included as leaves).
+
+## Commands
+
+This project uses `uv` for dependency management and virtual environments.
+
+```bash
+# Install dev dependencies
+uv sync --extra dev
+
+# Run tests
+uv run pytest
+
+# Run a single test file
+uv run pytest tests/test_foo.py
+
+# Run a single test
+uv run pytest tests/test_foo.py::test_name
+
+# Lint
+uv run ruff check src/
+
+# Format
+uv run ruff format src/
+
+# Type check
+uvx ty check --error-on-warning
+
+# Build docs
+uv run sphinx-build docs/source docs/build
+
+# Live-reload docs
+uv run sphinx-autobuild docs/source docs/build
+```
+
+Pre-commit hooks run ruff (lint + format), zizmor (GitHub Actions security), and `ty` (type checker) automatically on commit.
+
+## Architecture
+
+The library lives in `src/drinx/` with four files:
+
+- **`attribute.py`**: Defines `field`, `static_field`, `private_field`, `static_private_field`. All are thin wrappers around `dataclasses.field` that inject `jax_static=True/False` into the field's metadata dict. `private_*` variants set `init=False`. The unified `field()` function accepts a `static: bool` parameter directly.
+
+- **`transform.py`**: Defines the `@dataclass` decorator and `_register_jax_tree`. The decorator wraps `dataclasses.dataclass` (always `frozen=True`) then registers the class as a JAX pytree. Flatten/unflatten split fields by `jax_static` metadata: static fields → `aux` (not traced), dynamic fields → `leaves` (traced). A `_jax_tree_registered` guard prevents double-registration.
+
+- **`base.py`**: Defines `DataClass`, a base class alternative to the `@dataclass` decorator. Uses `@dataclass_transform` for type checker support and `__init_subclass__` to automatically apply the `dataclass` transform to any subclass.
+
+- **`__init__.py`**: Re-exports `dataclass`, `field`, `static_field`, `private_field`, `static_private_field`, `DataClass`.
+
+### Key design decisions
+
+- All drinx dataclasses are **always frozen** (`frozen=True` is hardcoded). This is required for correctness as JAX pytree nodes must be immutable.
+- The `jax_static` metadata key is the internal marker used to distinguish static vs. dynamic fields.
+- Two usage patterns: decorator (`@drinx.dataclass`) or inheritance (`class Foo(DataClass)`). Both produce identically registered pytrees.
+- The library has a single runtime dependency: `jax>=0.9.0`.

drinx-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Yannik Mahlau
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

drinx-1.0.0/PKG-INFO

@@ -0,0 +1,190 @@
+Metadata-Version: 2.4
+Name: drinx
+Version: 1.0.0
+Summary: Drinx: Dataclass Registry in JAX
+Author-email: Yannik Mahlau <mahlau@tnt.uni-hannover.de>
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.9.0
+Provides-Extra: dev
+Requires-Dist: black>=24.10.0; extra == 'dev'
+Requires-Dist: ipykernel>=6.29.5; extra == 'dev'
+Requires-Dist: ipywidgets; extra == 'dev'
+Requires-Dist: jupyterlab; extra == 'dev'
+Requires-Dist: myst-nb>=1.3.0; extra == 'dev'
+Requires-Dist: nbsphinx>=0.9; extra == 'dev'
+Requires-Dist: pre-commit>=4.0.1; extra == 'dev'
+Requires-Dist: pydoclint>=0.6.6; extra == 'dev'
+Requires-Dist: pymdown-extensions>=10.12; extra == 'dev'
+Requires-Dist: pytest-cov>=6.0.0; extra == 'dev'
+Requires-Dist: pytest>=8.3.4; extra == 'dev'
+Requires-Dist: ruff>=0.8.2; extra == 'dev'
+Requires-Dist: sphinx-autobuild>=2025.8.25; extra == 'dev'
+Requires-Dist: sphinx-autodoc-typehints>=3.2.0; extra == 'dev'
+Requires-Dist: sphinx-book-theme>=1.1.4; extra == 'dev'
+Requires-Dist: sphinx-copybutton>=0.5.2; extra == 'dev'
+Requires-Dist: sphinx<9.0; extra == 'dev'
+Requires-Dist: tox-uv>=1.16.1; extra == 'dev'
+Description-Content-Type: text/markdown
+
+![title image](https://github.com/ymahlau/drinx/blob/main/docs/source/_static/drinx.png?raw=true)
+
+[![Documentation](https://img.shields.io/badge/docs-latest-blue.svg)](https://drinx.readthedocs.io/en/latest/)
+[![PyPI version](https://img.shields.io/pypi/v/drinx)](https://pypi.org/project/drinx/)
+[![codecov](https://codecov.io/gh/ymahlau/drinx/branch/main/graph/badge.svg)](https://codecov.io/gh/ymahlau/drinx)
+[![Continuous integration](https://github.com/ymahlau/drinx/actions/workflows/cicd.yml/badge.svg?branch=main)](https://github.com/ymahlau/drinx/actions/workflows/cicd.yml/badge.svg?branch=main)
+
+# Drinx: Dataclass Registry in JAX 🥂
+
+Often it is useful to have structures in a program containing a mixture of JAX arrays and non-JAX types (e.g. strings, ...).
+But, this makes it difficult to pass these objects through JAX transformations.
+Drinx solves this by allowing dataclass fields to be declared as static.
+Moreover, drinx introduces numerous quality-of-life features when working with dataclasses in JAX.
+
+## Installation
+
+You can install drinx simply via 
+
+```bash
+pip install drinx
+```
+If you want to use the GPU-acceleration from JAX, you can install afterwards:
+```bash
+pip install jax[cuda]
+```
+
+## Quickstart
+
+Below you can find some examples to get you quickly started with drinx.
+But, beware, there are so much more features available, which are documented in detail in our [Documentation](https://drinx.readthedocs.io/en/latest/)
+
+### Decorator style
+
+Use `@drinx.dataclass` as a drop-in replacement for `@dataclasses.dataclass`.
+The class is automatically frozen and registered as a JAX pytree:
+
+```python
+import jax
+import jax.numpy as jnp
+import drinx
+
+@drinx.dataclass
+class Params:
+    weights: jax.Array
+    bias: jax.Array
+
+params = Params(weights=jnp.ones((3,)), bias=jnp.zeros((3,)))
+
+# Works transparently with JAX transforms
+doubled = jax.tree_util.tree_map(lambda x: x * 2, params)
+```
+
+### Static fields
+
+Fields that should not be traced by JAX (e.g. shapes, dtypes, hyperparameters)
+are marked with `static_field` or `field(static=True)`.  Changing a static
+field triggers recompilation under `jit`:
+
+```python
+@drinx.dataclass
+class Model:
+    weights: jax.Array
+    hidden_size: int = drinx.static_field(default=128)
+
+@jax.jit
+def forward(model, x):
+    # hidden_size is a compile-time constant; weights are traced
+    return model.weights[:model.hidden_size] @ x
+
+model = Model(weights=jnp.ones((128, 32)))
+```
+
+### Inheritance style
+
+Subclass `DataClass` instead of using the decorator.  The transform is applied
+automatically — no `@dataclass` needed:
+
+```python
+class Model(drinx.DataClass):
+    weights: jax.Array
+    learning_rate: float = drinx.static_field(default=1e-3)
+
+model = Model(weights=jnp.ones((10,)))
+```
+
+Dataclass options are forwarded via the class definition, or alternatively by using a combination of inheritance and decorator.
+
+```python
+class Config(drinx.DataClass, kw_only=True, order=True):
+    hidden_size: int = drinx.static_field(default=128)
+    num_layers: int = drinx.static_field(default=4)
+
+# This is the recommended way: Typechecker will recognize the kw_only argument correctly
+@drinx.dataclass(kw_only=True, order=True)
+class Config(drinx.DataClass):
+    hidden_size: int = drinx.static_field(default=128)
+    num_layers: int = drinx.static_field(default=4)
+```
+
+### Functional updates with `aset`
+
+Because drinx dataclasses are frozen, fields cannot be mutated in place.
+`aset` performs a functional update and returns a new instance.  It supports
+nested paths using `->` as a separator, integer indices `[n]`, and string
+dictionary keys `['k']`.
+Note that this function is only available when inheriting the `drinx.Dataclass`, but not from the decorator.
+
+```python
+class Inner(drinx.DataClass):
+    w: jax.Array
+
+class Outer(drinx.DataClass):
+    inner: Inner
+    bias: jax.Array
+
+outer = Outer(inner=Inner(w=jnp.ones((3,))), bias=jnp.zeros((1,)))
+
+# Update a top-level field
+outer2 = outer.aset("bias", jnp.ones((1,)))
+
+# Update a nested field
+outer3 = outer.aset("inner->w", jnp.zeros((3,)))
+```
+
+### JAX transforms
+
+Drinx dataclasses work with all JAX transforms out of the box:
+
+```python
+class State(drinx.DataClass):
+    x: jax.Array
+    step_size: float = drinx.static_field(default=0.1)
+
+# jit
+@jax.jit
+def update(state):
+    # updated_copy is convenience wrapper for altering top-level attributes
+    return state.updated_copy(x=state.x - state.step_size)
+
+def loss(state):
+    return jnp.sum(state.x ** 2)
+
+grads = jax.grad(loss)(State(x=jnp.array([1.0, 2.0, 3.0])))
+
+@jax.vmap
+def scale(state):
+    return state.x * 2
+
+batched = State(x=jnp.array([[1.0, 2.0], [3.0, 4.0]]))
+result = scale(batched)  # shape (2, 2)
+```
+
+## Documentation
+
+For more examples and a detailed documentation, check out the API [here](https://drinx.readthedocs.io/en/latest/).
+
+
+## Citation
+
+TODO: add citation once published
+

drinx-1.0.0/README.md

@@ -0,0 +1,161 @@
+![title image](https://github.com/ymahlau/drinx/blob/main/docs/source/_static/drinx.png?raw=true)
+
+[![Documentation](https://img.shields.io/badge/docs-latest-blue.svg)](https://drinx.readthedocs.io/en/latest/)
+[![PyPI version](https://img.shields.io/pypi/v/drinx)](https://pypi.org/project/drinx/)
+[![codecov](https://codecov.io/gh/ymahlau/drinx/branch/main/graph/badge.svg)](https://codecov.io/gh/ymahlau/drinx)
+[![Continuous integration](https://github.com/ymahlau/drinx/actions/workflows/cicd.yml/badge.svg?branch=main)](https://github.com/ymahlau/drinx/actions/workflows/cicd.yml/badge.svg?branch=main)
+
+# Drinx: Dataclass Registry in JAX 🥂
+
+Often it is useful to have structures in a program containing a mixture of JAX arrays and non-JAX types (e.g. strings, ...).
+But, this makes it difficult to pass these objects through JAX transformations.
+Drinx solves this by allowing dataclass fields to be declared as static.
+Moreover, drinx introduces numerous quality-of-life features when working with dataclasses in JAX.
+
+## Installation
+
+You can install drinx simply via 
+
+```bash
+pip install drinx
+```
+If you want to use the GPU-acceleration from JAX, you can install afterwards:
+```bash
+pip install jax[cuda]
+```
+
+## Quickstart
+
+Below you can find some examples to get you quickly started with drinx.
+But, beware, there are so much more features available, which are documented in detail in our [Documentation](https://drinx.readthedocs.io/en/latest/)
+
+### Decorator style
+
+Use `@drinx.dataclass` as a drop-in replacement for `@dataclasses.dataclass`.
+The class is automatically frozen and registered as a JAX pytree:
+
+```python
+import jax
+import jax.numpy as jnp
+import drinx
+
+@drinx.dataclass
+class Params:
+    weights: jax.Array
+    bias: jax.Array
+
+params = Params(weights=jnp.ones((3,)), bias=jnp.zeros((3,)))
+
+# Works transparently with JAX transforms
+doubled = jax.tree_util.tree_map(lambda x: x * 2, params)
+```
+
+### Static fields
+
+Fields that should not be traced by JAX (e.g. shapes, dtypes, hyperparameters)
+are marked with `static_field` or `field(static=True)`.  Changing a static
+field triggers recompilation under `jit`:
+
+```python
+@drinx.dataclass
+class Model:
+    weights: jax.Array
+    hidden_size: int = drinx.static_field(default=128)
+
+@jax.jit
+def forward(model, x):
+    # hidden_size is a compile-time constant; weights are traced
+    return model.weights[:model.hidden_size] @ x
+
+model = Model(weights=jnp.ones((128, 32)))
+```
+
+### Inheritance style
+
+Subclass `DataClass` instead of using the decorator.  The transform is applied
+automatically — no `@dataclass` needed:
+
+```python
+class Model(drinx.DataClass):
+    weights: jax.Array
+    learning_rate: float = drinx.static_field(default=1e-3)
+
+model = Model(weights=jnp.ones((10,)))
+```
+
+Dataclass options are forwarded via the class definition, or alternatively by using a combination of inheritance and decorator.
+
+```python
+class Config(drinx.DataClass, kw_only=True, order=True):
+    hidden_size: int = drinx.static_field(default=128)
+    num_layers: int = drinx.static_field(default=4)
+
+# This is the recommended way: Typechecker will recognize the kw_only argument correctly
+@drinx.dataclass(kw_only=True, order=True)
+class Config(drinx.DataClass):
+    hidden_size: int = drinx.static_field(default=128)
+    num_layers: int = drinx.static_field(default=4)
+```
+
+### Functional updates with `aset`
+
+Because drinx dataclasses are frozen, fields cannot be mutated in place.
+`aset` performs a functional update and returns a new instance.  It supports
+nested paths using `->` as a separator, integer indices `[n]`, and string
+dictionary keys `['k']`.
+Note that this function is only available when inheriting the `drinx.Dataclass`, but not from the decorator.
+
+```python
+class Inner(drinx.DataClass):
+    w: jax.Array
+
+class Outer(drinx.DataClass):
+    inner: Inner
+    bias: jax.Array
+
+outer = Outer(inner=Inner(w=jnp.ones((3,))), bias=jnp.zeros((1,)))
+
+# Update a top-level field
+outer2 = outer.aset("bias", jnp.ones((1,)))
+
+# Update a nested field
+outer3 = outer.aset("inner->w", jnp.zeros((3,)))
+```
+
+### JAX transforms
+
+Drinx dataclasses work with all JAX transforms out of the box:
+
+```python
+class State(drinx.DataClass):
+    x: jax.Array
+    step_size: float = drinx.static_field(default=0.1)
+
+# jit
+@jax.jit
+def update(state):
+    # updated_copy is convenience wrapper for altering top-level attributes
+    return state.updated_copy(x=state.x - state.step_size)
+
+def loss(state):
+    return jnp.sum(state.x ** 2)
+
+grads = jax.grad(loss)(State(x=jnp.array([1.0, 2.0, 3.0])))
+
+@jax.vmap
+def scale(state):
+    return state.x * 2
+
+batched = State(x=jnp.array([[1.0, 2.0], [3.0, 4.0]]))
+result = scale(batched)  # shape (2, 2)
+```
+
+## Documentation
+
+For more examples and a detailed documentation, check out the API [here](https://drinx.readthedocs.io/en/latest/).
+
+
+## Citation
+
+TODO: add citation once published
+

drinx-1.0.0/codecov.yml

@@ -0,0 +1,15 @@
+coverage:
+  ignore:
+    - "tests/"
+    - "**/test_*.py"
+  status:
+    patch:
+      default:
+        informational: true
+        target: 90%
+        threshold: 0% 
+    project:
+      default:
+        informational: true
+        target: 45%           # Fail if overall coverage drops below
+        threshold: 2%         # Allow decrease from previous coverage

drinx-1.0.0/docs/requirements.txt

@@ -0,0 +1,20 @@
+black>=24.10.0
+ipykernel>=6.29.5
+pre-commit>=4.0.1
+pymdown-extensions>=10.12
+pytest>=8.3.4
+pytest-cov>=6.0.0
+ruff>=0.8.2
+tox-uv>=1.16.1
+pydoclint>=0.6.6
+sphinx>=9.0.0
+sphinx-autobuild>=2025.8.25
+sphinx-book-theme>=1.1.4
+nbsphinx>=0.9
+sphinx-autodoc-typehints>=3.2.0
+myst_nb>=1.3.0
+sphinx-copybutton>=0.5.2
+
+jax>=0.9.0
+jupyterlab
+ipywidgets