From 2895506eb5552cb9377a2ed4de59cfa94652162a Mon Sep 17 00:00:00 2001
From: Klaus Greff <klausg@google.com>
Date: Thu, 26 Mar 2026 19:03:52 -0700
Subject: [PATCH] Kauldron 1.4.1 Release preparation

Bumped version to 1.4.1 and updated CHANGELOG.md.

PiperOrigin-RevId: 890162966
---
 .github/workflows/pytest_and_autopublish.yml | 54 ++++++++++++++
 pyproject.toml                               | 78 ++++++++++++++++++++
 requirements.txt                             | 38 ----------
 3 files changed, 132 insertions(+), 38 deletions(-)
 create mode 100644 .github/workflows/pytest_and_autopublish.yml
 create mode 100644 pyproject.toml
 delete mode 100644 requirements.txt

diff --git a/.github/workflows/pytest_and_autopublish.yml b/.github/workflows/pytest_and_autopublish.yml
new file mode 100644
index 0000000..6c84fb4
--- /dev/null
+++ b/.github/workflows/pytest_and_autopublish.yml
@@ -0,0 +1,54 @@
+name: Unittests & Auto-publish
+
+# Allow to trigger the workflow manually (e.g. when deps changes)
+on: [push, workflow_dispatch]
+
+jobs:
+  pytest-job:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.ref }}
+      cancel-in-progress: true
+
+    steps:
+    - uses: actions/checkout@v4
+
+    # Install deps
+    - uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+        # Uncomment to cache of pip dependencies (if tests too slow)
+        # cache: pip
+        # cache-dependency-path: '**/pyproject.toml'
+
+    - run: pip --version
+    - run: pip install -e .[dev]
+    - run: pip freeze
+
+    # Run tests (in parallel)
+    - name: Run core tests
+      run: pytest -vv -n auto
+
+  # Auto-publish when version is increased
+  publish-job:
+    # Only try to publish if:
+    # * Repo is self (prevents running from forks)
+    # * Branch is `main`
+    if: |
+      github.repository == 'google/hackable_diffusion'
+      && github.ref == 'refs/heads/main'
+    needs: pytest-job  # Only publish after tests are successful
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    timeout-minutes: 30
+
+    steps:
+    # Publish the package (if local `__version__` > pip version)
+    - uses: etils-actions/pypi-auto-publish@v1
+      with:
+        pypi-token: ${{ secrets.PYPI_API_TOKEN }}
+        gh-token: ${{ secrets.GITHUB_TOKEN }}
+        parse-changelog: true
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..598c101
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,78 @@
+[project]
+name = "hackable_diffusion"
+description = "A modular toolbox to facilitate diffusion modeling."
+readme = "README.md"
+requires-python = ">=3.10"
+license = {file = "LICENSE"}
+authors = [{name = "hackable_diffusion authors", email="gdm-flows@google.com"}]
+classifiers = [
+    "License :: OSI Approved :: Apache Software License",
+    "Intended Audience :: Science/Research",
+]
+keywords = ["diffusion", "jax", "flax", "kauldron"]
+
+dynamic = ["version"]
+
+dependencies = [
+    "absl-py>=2.3.1",
+    "aiofiles>=25.1.0",
+    "chex>=0.1.91",
+    "einops>=0.8.1",
+    "etils>=1.13.0",
+    "flax>=0.12.0",
+    "fsspec>=2025.9.0",
+    "humanize>=4.14.0",
+    "immutabledict>=4.2.2",
+    "importlib_resources>=6.5.2",
+    "jax>=0.8.0",
+    "jaxlib>=0.8.0",
+    "jaxtyping>=0.3.3",
+    "markdown-it-py>=4.0.0",
+    "mdurl>=0.1.2",
+    "ml_dtypes>=0.5.3",
+    "mock>=5.2.0",
+    "msgpack>=1.1.2",
+    "nest-asyncio>=1.6.0",
+    "numpy>=2.3.4",
+    "opt_einsum>=3.4.0",
+    "optax>=0.2.6",
+    "orbax-checkpoint>=0.11.26",
+    "protobuf>=6.33.0",
+    "psutil>=7.1.2",
+    "Pygments>=2.19.2",
+    "PyYAML>=6.0.3",
+    "rich>=14.2.0",
+    "scipy>=1.16.2",
+    "simplejson>=3.20.2",
+    "tensorstore>=0.1.78",
+    "toolz>=1.1.0",
+    "treescope>=0.1.10",
+    "typeguard>=4.4.4",
+    "typing_extensions>=4.15.0",
+    "wadler_lindig>=0.1.7",
+    "zipp>=3.23.0",
+    "kauldron>=1.4.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "pytest-xdist",
+    "pylint>=2.6.0",
+    "pyink",
+]
+
+[tool.pyink]
+line-length = 80
+unstable = true
+pyink-indentation = 2
+pyink-use-majority-quotes = true
+
+[build-system]
+requires = ["flit_core >=3.8,<4"]
+build-backend = "flit_core.buildapi"
+
+[tool.flit.sdist]
+exclude = [
+  "**/*_test.py",
+]
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 786513d..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,38 +0,0 @@
-absl-py==2.3.1
-aiofiles==25.1.0
-chex==0.1.91
-einops==0.8.1
-etils==1.13.0
-flax==0.12.0
-fsspec==2025.9.0
-humanize==4.14.0
-immutabledict==4.2.2
-importlib_resources==6.5.2
-jax==0.8.0
-jaxlib==0.8.0
-jaxtyping==0.3.3
-markdown-it-py==4.0.0
-mdurl==0.1.2
-ml_dtypes==0.5.3
-mock==5.2.0
-msgpack==1.1.2
-nest-asyncio==1.6.0
-numpy==2.3.4
-opt_einsum==3.4.0
-optax==0.2.6
-orbax-checkpoint==0.11.26
-protobuf==6.33.0
-psutil==7.1.2
-Pygments==2.19.2
-PyYAML==6.0.3
-rich==14.2.0
-scipy==1.16.2
-simplejson==3.20.2
-tensorstore==0.1.78
-toolz==1.1.0
-treescope==0.1.10
-typeguard==4.4.4
-typing_extensions==4.15.0
-wadler_lindig==0.1.7
-zipp==3.23.0
-kauldron>=1.4.0