[nnx] add flaxlib

.github/workflows/pythonpublish.yml → .github/workflows/flax_publish.yml

@@ -1,11 +1,11 @@
 # This workflows will upload a Python Package using Twine when a release is created
 # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
 
-name: Upload Python Package
+name: Flax - Build and upload to PyPI
 
 on:
   release:
-    types: [created]
+    types: [published]
 
 jobs:
   deploy:

.github/workflows/build.yml → .github/workflows/flax_test.yml

@@ -1,7 +1,7 @@
 # This workflow will install Python dependencies, run tests and lint with a variety of Python versions
 # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
 
-name: Build
+name: Flax - Test
 
 on:
   push:
@@ -70,7 +70,7 @@ jobs:
       uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
-    - uses: yezz123/setup-uv@v4
+    - uses: astral-sh/setup-uv@v2
       with:
         uv-version: "0.3.0"
     - name: Install standalone dependencies only
@@ -104,23 +104,17 @@ jobs:
       uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
-    - uses: yezz123/setup-uv@v4
+    - name: Setup uv
+      uses: astral-sh/setup-uv@v2
       with:
-        uv-version: "0.3.0"
-    - name: Cached virtual environment
-      id: venv_cache
-      uses: actions/cache@v3
-      with:
-        path: .venv
-        key: pip-${{ steps.setup_python.outputs.python-version }}-${{ hashFiles('uv.lock') }}
-    - name: Install Dependencies for cache
-      if: steps.venv_cache.outputs.cache-hit != 'true'
-      run: |
-        if [ -d ".venv" ]; then rm -rf .venv; fi
-        uv sync --locked --all-extras
-    - name: Check lockfile
+        version: "0.3.0"
+    - name: Setup Rust (flaxlib)
+      uses: actions-rust-lang/setup-rust-toolchain@v1
+
+    - name: Install dependencies
       run: |
-        uv sync --locked --all-extras
+        uv sync --locked --extra all --extra testing --extra docs
+        uv pip install ./flaxlib
     - name: Install JAX
       run: |
         if [[ "${{ matrix.jax-version }}" == "newest" ]]; then

.github/workflows/flaxlib_publish.yml

@@ -0,0 +1,94 @@
+name: Flaxlib - Build and upload to PyPI
+
+# for testing only:
+on:
+  push:
+    branches: [main]
+    paths: ['flaxlib/**']
+  release:
+    types: [published]
+
+jobs:
+  build_wheels:
+    if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')
+    name: Build wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        # macos-13 is an intel runner, macos-14 is apple silicon
+        os: [ubuntu-latest, windows-latest, macos-13, macos-14]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+
+    - name: Setup Rust
+      uses: actions-rust-lang/setup-rust-toolchain@v1
+
+    - name: Install cibuildwheel
+      run: python -m pip install cibuildwheel==2.21.0
+
+    - name: Build wheels
+      run: python -m cibuildwheel --output-dir ./flaxlib/wheelhouse ./flaxlib
+      env:
+        # rust doesn't seem to be available for musl linux on i686
+        CIBW_SKIP: "*-musllinux_i686"
+        CIBW_ENVIRONMENT: 'PATH="$HOME/.cargo/bin:$PATH" CARGO_TERM_COLOR="always"'
+        CIBW_ENVIRONMENT_WINDOWS: 'PATH="$UserProfile\.cargo\bin;$PATH"'
+        CIBW_BEFORE_BUILD: rustup show
+        CIBW_BEFORE_BUILD_LINUX: |
+          curl https://sh.rustup.rs -sSf | sh -s -- --default-toolchain=stable --profile=minimal -y &&
+          rustup show
+
+    - uses: actions/upload-artifact@v4
+      with:
+        name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
+        path: ./flaxlib/wheelhouse/*.whl
+
+  build_sdist:
+    if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')
+    name: Build source distribution
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Setup Rust
+      uses: actions-rust-lang/setup-rust-toolchain@v1
+
+    - name: Build sdist
+      run: pipx run build --sdist flaxlib
+
+    - uses: actions/upload-artifact@v4
+      with:
+        name: cibw-sdist
+        path: ./flaxlib/dist/*.tar.gz
+
+  upload_pypi:
+    if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')
+    name: Upload to PyPI
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    steps:
+    - uses: actions/setup-python@v1
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools build wheel twine
+    - uses: actions/download-artifact@v4
+      with:
+        # unpacks all CIBW artifacts into dist/
+        pattern: cibw-*
+        path: ./flaxlib/dist
+        merge-multiple: true
+
+    - name: Build and publish
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.FLAXLIB_PYPI_TOKEN }}
+      run: |
+        twine upload flaxlib/dist/*

flax/nnx/graph.py

@@ -51,7 +51,7 @@
 AuxData = tp.TypeVar('AuxData')
 
 StateLeaf = VariableState[tp.Any]
-NodeLeaf = VariableState[tp.Any]
+NodeLeaf = Variable[tp.Any]
 GraphState = State[Key, StateLeaf]
 GraphFlatState = FlatState[StateLeaf]
 

flaxlib/.gitignore

@@ -0,0 +1,75 @@
+/target
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+.pytest_cache/
+*.py[cod]
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+.venv/
+env/
+bin/
+build/
+develop-eggs/
+dist/
+eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+include/
+man/
+venv/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+pip-selfcheck.json
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.cache
+nosetests.xml
+coverage.xml
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+
+# Rope
+.ropeproject
+
+# Django stuff:
+*.log
+*.pot
+
+.DS_Store
+
+# Sphinx documentation
+docs/_build/
+
+# PyCharm
+.idea/
+
+# VSCode
+.vscode/
+
+# Pyenv
+.python-version
+
+# cibuildwheel
+/wheelhouse