Add lock around init func and prepare release

.github/workflows/ci.yml

@@ -1,8 +1,3 @@
-# This file is autogenerated by maturin v1.8.2
-# To update, run
-#
-#    maturin generate-ci --pytest --platform all --zig github
-#
 name: CI
 
 on:
@@ -31,16 +26,32 @@ jobs:
             target: aarch64
     steps:
       - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
       - name: Install uv
         uses: astral-sh/setup-uv@v5
+      - uses: actions/setup-python@v5
+        with:
+          python-version: |
+            3.10
+            3.11
+            3.12
+      # Leave out 3.13 on aarch due to an issue in pyo3/rust-numpy 0.23.4
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        if: ${{ matrix.platform.target == 'aarch64' }}
+        with:
+          target: ${{ matrix.platform.target }}
+          args: --release --out dist --interpreter 3.10 3.11 3.12 --zig
+          sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
+          manylinux: auto
+          before-script-linux: |
+            dnf install -y clang-libs clang || sudo apt install llvm-dev libclang-dev clang
       - name: Build wheels
         uses: PyO3/maturin-action@v1
+        if: ${{ matrix.platform.target == 'x86_64' }}
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --find-interpreter --zig
+          # No py3.13 yet...
+          args: --release --out dist --interpreter 3.10 3.11 3.12 --zig
           sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
           manylinux: auto
           before-script-linux: |
@@ -160,7 +171,11 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: |
+            3.10
+            3.11
+            3.12
+          # 3.13 leave out 3.13 due to a segfault
           architecture: ${{ matrix.platform.target }}
       - name: Install uv
         uses: astral-sh/setup-uv@v5
@@ -211,7 +226,10 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: |
+            3.10
+            3.11
+            3.12
       - name: Install uv
         uses: astral-sh/setup-uv@v5
       - uses: maxim-lobanov/setup-xcode@v1

CHANGELOG.md

@@ -2,7 +2,19 @@
 
 All notable changes to this project will be documented in this file.
 
-## [0.13.3] - 2025-02-11
+## [0.13.4] - 2025-02-18
+
+### Bug Fixes
+
+- Add lock for pymc init point func (Adrian Seyboldt)
+
+
+### Ci
+
+- Make sure all python versions are available in the builds (Adrian Seyboldt)
+
+
+## [0.13.3] - 2025-02-12
 
 ### Bug Fixes
 
@@ -44,6 +56,18 @@ All notable changes to this project will be documented in this file.
 - Reformat some code (Adrian Seyboldt)
 
 
+### Build
+
+- Bump some dependency versions (Adrian Seyboldt)
+
+
+### Ci
+
+- Use ubuntu_latest on aarch64 (Adrian Seyboldt)
+
+- Update CI script using maturin (Adrian Seyboldt)
+
+
 ## [0.13.2] - 2024-07-26
 
 ### Features
@@ -178,8 +202,6 @@ All notable changes to this project will be documented in this file.
 
 ### Ci
 
-- Fix uploads of releases (Adrian Seyboldt)
-
 - Fix architectures in CI (Adrian Seyboldt)
 
 
@@ -227,6 +249,11 @@ All notable changes to this project will be documented in this file.
 - Set the number of parallel chains dynamically (Adrian Seyboldt)
 
 
+### Ci
+
+- Fix uploads of releases (Adrian Seyboldt)
+
+
 ## [0.9.2] - 2024-02-19
 
 ### Bug Fixes

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nutpie"
-version = "0.13.3"
+version = "0.13.4"
 authors = [
   "Adrian Seyboldt <[email protected]>",
   "PyMC Developers <[email protected]>",

python/nutpie/compile_pymc.py

@@ -1,5 +1,6 @@
 import dataclasses
 import itertools
+import threading
 import warnings
 from collections.abc import Iterable
 from dataclasses import dataclass
@@ -31,7 +32,7 @@ def intrinsic(f):
     from pytensor.tensor import TensorVariable, Variable
 
 
-def rv_dict_to_flat_array_wrapper(
+def _rv_dict_to_flat_array_wrapper(
     fn: Callable[[SeedType], dict[str, np.ndarray]],
     names: list[str],
     shapes: list[tuple[int]],
@@ -509,6 +510,8 @@ def compile_pymc_model(
         return_transformed=True,
     )
 
+    initial_point_fn = _wrap_with_lock(initial_point_fn)
+
     if backend.lower() == "numba":
         if gradient_backend == "jax":
             raise ValueError("Gradient backend cannot be jax when using numba backend")
@@ -530,7 +533,18 @@ def compile_pymc_model(
         raise ValueError(f"Backend must be one of numba and jax. Got {backend}")
 
 
-def _compute_shapes(model):
+def _wrap_with_lock(func: Callable) -> Callable:
+    lock = threading.Lock()
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        with lock:
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+def _compute_shapes(model) -> dict[str, tuple[int, ...]]:
     import pytensor
     from pymc.initial_point import make_initial_point_fn
 
@@ -663,7 +677,7 @@ def _make_functions(
 
     num_free_vars = count
 
-    initial_point_fn = rv_dict_to_flat_array_wrapper(
+    initial_point_fn = _rv_dict_to_flat_array_wrapper(
         pymc_initial_point_fn, names=joined_names, shapes=joined_shapes
     )