From f210a3c16fc96d7b1c21fd2cbada30373564631d Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Wed, 13 May 2026 23:07:10 +0200 Subject: [PATCH 1/4] Activate beartype claw on lcm.regime_building MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `lcm.regime_building` to the import-claw registrations in `lcm/__init__.py`, mapping type violations to `ModelInitializationError` (regime compilation is part of model construction). Wraps `dags.with_signature` / `dags.rename_arguments` via a thin helper in `lcm.utils._dags_forwarders` that defaults `forwarder=True`. Each direct caller of those dags helpers inside pylcm produces a generic `*args, **kwargs` forwarder whose annotations describe the inner function's contract, not the wrapper's own call protocol. `forwarder=True` advertises the wrapper as a permissive forwarder on its `__annotations__`, so beartype's claw treats it as universally permissive and skips per-parameter enforcement — matching the wrapper's actual runtime behaviour. dags' own `get_annotations` recovers the user-described view via its existing args/kwargs-mismatch fallback. Pins `dags` to the `feat/no-type-check-flag` branch (PR OpenSourceEconomics/dags#82) which adds the `forwarder` flag. Will be replaced with a released version once that PR lands. Annotation drift fixed alongside activation: - `collect_state_transitions(states: ...)` widened to `Mapping[StateName, Grid | None]`; test mocks pass `None` for placeholder states. - `map_coordinates(coordinates: ...)` widened to `Sequence[Array] | Array`; callers pass a 2D `jnp.array` (a single Array, not a sequence) and JAX produces a single tracer under vmap. - `_get_weights_func_for_shock.weights_func_runtime.shock_kw` typed as `dict[str, float | FloatND]`; under JIT the runtime shock params arrive as tracers. - `solve_brute.solve.running_any_nan` / `running_any_inf` typed as `BoolND` to match the underlying `jnp.zeros((), dtype=bool)`. - `diagnostics._wrap_with_reduction.reduced.**kwargs` typed as `Array | Mapping[str, Array]` since `next_regime_to_V_arr` flows through it as a mapping alongside Array-valued state/action inputs. - One test fixture (`tests/test_next_state.py::test_create_stochastic_next_func`) updated to pass an `int32`-typed `labels` array. Out of scope: extending the claw to `lcm.solution` / `lcm.simulation`, which surfaces further annotation drift the claw correctly catches but that needs its own pass. Co-Authored-By: Claude Opus 4.7 --- pixi.lock | 47 ++++++------ pyproject.toml | 5 ++ src/lcm/__init__.py | 3 +- src/lcm/_beartype_conf.py | 4 + src/lcm/regime_building/Q_and_F.py | 3 +- src/lcm/regime_building/V.py | 3 +- src/lcm/regime_building/diagnostics.py | 7 +- src/lcm/regime_building/max_Q_over_a.py | 2 +- src/lcm/regime_building/ndimage.py | 2 +- src/lcm/regime_building/next_state.py | 3 +- src/lcm/regime_building/processing.py | 11 ++- src/lcm/regime_building/validation.py | 2 +- src/lcm/solution/solve_brute.py | 10 +-- src/lcm/utils/_dags_forwarders.py | 98 +++++++++++++++++++++++++ tests/test_next_state.py | 2 +- 15 files changed, 160 insertions(+), 42 deletions(-) create mode 100644 src/lcm/utils/_dags_forwarders.py diff --git a/pixi.lock b/pixi.lock index 05c13a7af..3b745fb11 100644 --- a/pixi.lock +++ b/pixi.lock @@ -275,7 +275,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2c/1a/aff8bb287a4b1400f69e09a53bd65de96aa5cee5691925b38731c67fc695/click_default_group-1.2.4-py2.py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/8d/2d/f61c918d9edc2127068f0d5ad4604fedd9bfd393f464219090f3279c73f7/estimagic-0.5.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d5/1f/5f4a3cd9e4440e9d9bc78ad0a91a1c8d46b4d429d5239ebe6793c9fe5c41/fsspec-2026.3.0-py3-none-any.whl @@ -588,7 +588,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/78/a3d9ceda0793f4fb43daa292af7b801932611a1aed442636ddfc93d58c7a/jax_cuda12_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -888,7 +888,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/21/98/77f15d81fd0637da454e453c8456d4a2b5c8b2e66823b4237ee8689152cf/jax_cuda13_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -1158,7 +1158,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -1399,7 +1399,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -1637,7 +1637,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-1.3.2-hfd05255_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -1895,7 +1895,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -2146,7 +2146,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -2393,7 +2393,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-1.3.2-hfd05255_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -2648,7 +2648,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl @@ -2911,7 +2911,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -3166,7 +3166,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -3417,7 +3417,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-1.3.2-hfd05255_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -3716,7 +3716,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/78/a3d9ceda0793f4fb43daa292af7b801932611a1aed442636ddfc93d58c7a/jax_cuda12_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -4029,7 +4029,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/21/98/77f15d81fd0637da454e453c8456d4a2b5c8b2e66823b4237ee8689152cf/jax_cuda13_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -4308,7 +4308,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl @@ -4673,7 +4673,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-ng-2.3.3-hceb46e0_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -4964,7 +4964,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-ng-2.3.3-hed4e4f5_1.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -5274,7 +5274,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-ng-2.3.3-h0261ad2_1.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -7542,10 +7542,9 @@ packages: purls: [] size: 210103 timestamp: 1771943128249 -- pypi: https://files.pythonhosted.org/packages/2c/c1/a662f0a8f6e024fca239d493f278d9adf5de1c8408af46a53a76beb13534/dags-0.5.1-py3-none-any.whl +- pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 name: dags - version: 0.5.1 - sha256: e9fd9fbe0536784fe8b8ce58ea194801b1de39d7364941d4a1f2d8240c14123d + version: 0.5.2.dev1+gd415d10ee requires_dist: - flatten-dict - networkx>=3.6 @@ -14086,8 +14085,8 @@ packages: timestamp: 1774796815820 - pypi: ./ name: pylcm - version: 0.0.2.dev136+ga5d932a48.d20260513 - sha256: fc3e3cff622b9db6e1f5a01c6cfad41879d5d1efdb982743f18b8ca3edf585a7 + version: 0.0.2.dev179+g59aca3f5f.d20260513 + sha256: 0c5af9810a5be3f1e5861656eaf8e0c6fa09dcbf505d6e746feffeff83e0cce6 requires_dist: - beartype>=0.21 - cloudpickle>=3.1.2 diff --git a/pyproject.toml b/pyproject.toml index 3651a96a5..e23e07d64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -149,6 +149,11 @@ ty = "ty check" jax = ">=0.9" pdbp = "*" pylcm = { path = ".", editable = true } +# Pin dags to the feat/no-type-check-flag branch (PR +# OpenSourceEconomics/dags#82): adds the `forwarder=True` flag we need +# to make `with_signature` cooperate with beartype's import claw. +# Will be replaced with a released version once that PR lands. +dags = { git = "https://github.com/OpenSourceEconomics/dags.git", rev = "d415d10" } [tool.pixi.tasks] asv-compare = "asv compare" asv-preview = "asv preview" diff --git a/src/lcm/__init__.py b/src/lcm/__init__.py index bb00b6127..565acfce4 100644 --- a/src/lcm/__init__.py +++ b/src/lcm/__init__.py @@ -38,11 +38,12 @@ # exception most natural to that subpackage (see `lcm._beartype_conf`). from beartype.claw import beartype_package -from lcm._beartype_conf import GRID_CONF, PARAMS_CONF +from lcm._beartype_conf import GRID_CONF, PARAMS_CONF, REGIME_BUILDING_CONF beartype_package("lcm.grids", conf=GRID_CONF) beartype_package("lcm.shocks", conf=GRID_CONF) beartype_package("lcm.params", conf=PARAMS_CONF) +beartype_package("lcm.regime_building", conf=REGIME_BUILDING_CONF) from lcm import shocks # noqa: E402 from lcm._version import __version__ # noqa: E402 diff --git a/src/lcm/_beartype_conf.py b/src/lcm/_beartype_conf.py index daf276120..6a7bee7cf 100644 --- a/src/lcm/_beartype_conf.py +++ b/src/lcm/_beartype_conf.py @@ -66,3 +66,7 @@ def deco(cls: type[C]) -> type[C]: # Used on `Model.solve` and `Model.simulate`. PARAMS_CONF = _conf(InvalidParamsError) + +# Used by the claw on `lcm.regime_building` (regime compilation pipeline, +# part of model construction). +REGIME_BUILDING_CONF = _conf(ModelInitializationError) diff --git a/src/lcm/regime_building/Q_and_F.py b/src/lcm/regime_building/Q_and_F.py index 814619ddf..9bff01a85 100644 --- a/src/lcm/regime_building/Q_and_F.py +++ b/src/lcm/regime_building/Q_and_F.py @@ -3,7 +3,7 @@ from typing import Any, cast import jax.numpy as jnp -from dags import concatenate_functions, with_signature +from dags import concatenate_functions from jax import Array from lcm.regime_building.h_dag import _get_build_H_kwargs @@ -24,6 +24,7 @@ TransitionFunctionName, TransitionFunctionsMapping, ) +from lcm.utils._dags_forwarders import with_signature from lcm.utils.dispatchers import productmap from lcm.utils.functools import get_union_of_args diff --git a/src/lcm/regime_building/V.py b/src/lcm/regime_building/V.py index 31521a5c8..ad8c7d21a 100644 --- a/src/lcm/regime_building/V.py +++ b/src/lcm/regime_building/V.py @@ -3,7 +3,7 @@ from types import MappingProxyType import jax.numpy as jnp -from dags import concatenate_functions, with_signature +from dags import concatenate_functions from dags.tree import qname_from_tree_path from jax import Array @@ -13,6 +13,7 @@ from lcm.regime_building.ndimage import map_coordinates from lcm.shocks import _ShockGrid from lcm.typing import FloatND, ScalarFloat, StateName +from lcm.utils._dags_forwarders import with_signature from lcm.utils.functools import all_as_kwargs from lcm.variables import Variables, get_grids diff --git a/src/lcm/regime_building/diagnostics.py b/src/lcm/regime_building/diagnostics.py index fb0ac9f17..ad88e6cce 100644 --- a/src/lcm/regime_building/diagnostics.py +++ b/src/lcm/regime_building/diagnostics.py @@ -9,7 +9,7 @@ `lcm.utils.error_handling`. """ -from collections.abc import Callable +from collections.abc import Callable, Mapping from types import MappingProxyType from typing import Any @@ -164,7 +164,10 @@ def _wrap_with_reduction( """ - def reduced(**kwargs: Array) -> dict[str, Any]: + # `kwargs` carries the wrapped function's full input map, which may + # include the `next_regime_to_V_arr` mapping alongside the Array-valued + # state/action inputs — so the kwarg-value type is `Array | Mapping`. + def reduced(**kwargs: Array | Mapping[str, Array]) -> dict[str, Any]: U_arr, F_arr, E_next_V, Q_arr, regime_probs = func(**kwargs) F_float = F_arr.astype(float) # NaN-count arrays are masked by feasibility: only feasible cells diff --git a/src/lcm/regime_building/max_Q_over_a.py b/src/lcm/regime_building/max_Q_over_a.py index f1e5d3913..d5fa18b93 100644 --- a/src/lcm/regime_building/max_Q_over_a.py +++ b/src/lcm/regime_building/max_Q_over_a.py @@ -3,7 +3,6 @@ from types import MappingProxyType import jax.numpy as jnp -from dags import with_signature from jax import Array from lcm.regime_building.argmax import argmax_and_max @@ -17,6 +16,7 @@ RegimeName, StateName, ) +from lcm.utils._dags_forwarders import with_signature from lcm.utils.dispatchers import productmap diff --git a/src/lcm/regime_building/ndimage.py b/src/lcm/regime_building/ndimage.py index 8d7d80d97..69a812b02 100644 --- a/src/lcm/regime_building/ndimage.py +++ b/src/lcm/regime_building/ndimage.py @@ -25,7 +25,7 @@ @jit def map_coordinates( input: Array, # noqa: A002 - coordinates: Sequence[Array], + coordinates: Sequence[Array] | Array, ) -> Array: """Map the input array to new coordinates using linear interpolation. diff --git a/src/lcm/regime_building/next_state.py b/src/lcm/regime_building/next_state.py index 1e81f0f23..42fffc1b4 100644 --- a/src/lcm/regime_building/next_state.py +++ b/src/lcm/regime_building/next_state.py @@ -4,7 +4,7 @@ from types import MappingProxyType import jax -from dags import concatenate_functions, with_signature +from dags import concatenate_functions from dags.tree import qname_from_tree_path from jax import Array @@ -26,6 +26,7 @@ TransitionFunctionName, TransitionFunctionsMapping, ) +from lcm.utils._dags_forwarders import with_signature from lcm.variables import Variables diff --git a/src/lcm/regime_building/processing.py b/src/lcm/regime_building/processing.py index 0b0c78945..8f19ca1c8 100644 --- a/src/lcm/regime_building/processing.py +++ b/src/lcm/regime_building/processing.py @@ -6,8 +6,7 @@ from typing import Any, Literal, cast import jax -from dags import concatenate_functions, get_annotations, with_signature -from dags.signature import rename_arguments +from dags import concatenate_functions, get_annotations from dags.tree import QNAME_DELIMITER, qname_from_tree_path, tree_path_from_qname from jax import Array from jax import numpy as jnp @@ -45,6 +44,7 @@ from lcm.typing import ( ArgmaxQOverAFunction, Float1D, + FloatND, FunctionsMapping, Int1D, InternalUserFunction, @@ -63,6 +63,7 @@ UserFunction, VmappedRegimeTransitionFunction, ) +from lcm.utils._dags_forwarders import rename_arguments, with_signature from lcm.utils.containers import ensure_containers_are_immutable from lcm.utils.dispatchers import simulation_spacemap, vmap_1d from lcm.utils.namespace import flatten_regime_namespace, unflatten_regime_namespace @@ -888,7 +889,11 @@ def _get_weights_func_for_shock(*, name: str, grid: _ShockGrid) -> UserFunction: @with_signature(args=args, return_annotation="FloatND", enforce=False) def weights_func_runtime(*a: Array, **kwargs: Array) -> Float1D: # noqa: ARG001 - shock_kw: dict[str, float] = { # ty: ignore[invalid-assignment] + # `float` here covers Python floats from fixed_params; under + # JIT tracing, the runtime values forwarded through `kwargs` + # arrive as JAX tracers (`FloatND`), which are accepted by the + # shock grid's `compute_gridpoints` / `compute_transition_probs`. + shock_kw: dict[str, float | FloatND] = { **fixed_params, **{raw: kwargs[qn] for qn, raw in runtime_param_names.items()}, } diff --git a/src/lcm/regime_building/validation.py b/src/lcm/regime_building/validation.py index 5457e03f5..471d16548 100644 --- a/src/lcm/regime_building/validation.py +++ b/src/lcm/regime_building/validation.py @@ -241,7 +241,7 @@ def _find_function_output_grid_indexing( def collect_state_transitions( - states: Mapping[StateName, Grid], + states: Mapping[StateName, Grid | None], state_transitions: Mapping[ StateName, UserFunction | Callable | None | Mapping[RegimeName, UserFunction | Callable], diff --git a/src/lcm/solution/solve_brute.py b/src/lcm/solution/solve_brute.py index 57b0c9306..a87eb9ffc 100644 --- a/src/lcm/solution/solve_brute.py +++ b/src/lcm/solution/solve_brute.py @@ -12,7 +12,7 @@ from lcm.ages import AgeGrid from lcm.interfaces import InternalRegime, _build_regime_sharding -from lcm.typing import FloatND, InternalParams, RegimeName, StateName +from lcm.typing import BoolND, FloatND, InternalParams, RegimeName, StateName from lcm.utils.error_handling import validate_V from lcm.utils.logging import ( format_duration, @@ -107,8 +107,8 @@ def solve( diagnostic_min: list[FloatND] = [] diagnostic_max: list[FloatND] = [] diagnostic_mean: list[FloatND] = [] - running_any_nan: FloatND = jnp.zeros((), dtype=bool) - running_any_inf: FloatND = jnp.zeros((), dtype=bool) + running_any_nan: BoolND = jnp.zeros((), dtype=bool) + running_any_inf: BoolND = jnp.zeros((), dtype=bool) logger.info("Starting solution") total_start = time.monotonic() @@ -471,8 +471,8 @@ def _emit_post_loop_diagnostics( solution: MappingProxyType[int, MappingProxyType[RegimeName, FloatND]], internal_regimes: Mapping[RegimeName, InternalRegime], internal_params: InternalParams, - running_any_nan: FloatND, - running_any_inf: FloatND, + running_any_nan: BoolND, + running_any_inf: BoolND, diagnostic_min: list[FloatND] | None, diagnostic_max: list[FloatND] | None, diagnostic_mean: list[FloatND] | None, diff --git a/src/lcm/utils/_dags_forwarders.py b/src/lcm/utils/_dags_forwarders.py new file mode 100644 index 000000000..7ba1ab0cb --- /dev/null +++ b/src/lcm/utils/_dags_forwarders.py @@ -0,0 +1,98 @@ +"""Thin wrappers around `dags.signature` that default to `forwarder=True`. + +Every direct caller of `dags.with_signature` / `dags.rename_arguments` +inside pylcm produces a generic `*args, **kwargs` forwarder whose +annotations describe the inner function's contract, not the wrapper's +own runtime call protocol. Setting `forwarder=True` on each call tells +beartype's import claw to treat the wrapper as permissive — matching +the wrapper's actual behaviour — instead of enforcing user annotations +against JAX tracers and other forwarded objects. + +Importing these wrappers (rather than the dags ones) keeps that intent +visible at each call site and concentrates the configuration in one +place if pylcm ever needs to flip the default. + +""" + +import inspect +from collections.abc import Callable, Mapping, Sequence +from typing import Any, overload + +from dags import rename_arguments as _rename_arguments +from dags import with_signature as _with_signature +from dags.typing import P, R + + +@overload +def with_signature( + func: Callable[P, R], + *, + args: Mapping[str, str] | Sequence[str] | None = None, + kwargs: Mapping[str, str] | Sequence[str] | None = None, + enforce: bool = True, + return_annotation: Any = inspect.Parameter.empty, # noqa: ANN401 +) -> Callable[P, R]: ... + + +@overload +def with_signature( + *, + args: Mapping[str, str] | Sequence[str] | None = None, + kwargs: Mapping[str, str] | Sequence[str] | None = None, + enforce: bool = True, + return_annotation: Any = inspect.Parameter.empty, # noqa: ANN401 +) -> Callable[[Callable[P, R]], Callable[P, R]]: ... + + +def with_signature( + func: Callable[P, R] | None = None, + *, + args: Mapping[str, str] | Sequence[str] | None = None, + kwargs: Mapping[str, str] | Sequence[str] | None = None, + enforce: bool = True, + return_annotation: Any = inspect.Parameter.empty, +) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]: + """`dags.with_signature` with `forwarder=True` baked in. + + See `dags.signature.with_signature` for argument semantics. The + `forwarder=True` setting advertises the resulting wrapper as a + permissive `*args, **kwargs` forwarder on `__annotations__`, + keeping beartype's claw from enforcing per-parameter annotations + against the forwarder's actual call arguments. + """ + return _with_signature( + func, # ty: ignore[invalid-argument-type] + args=args, + kwargs=kwargs, + enforce=enforce, + return_annotation=return_annotation, + forwarder=True, + ) + + +@overload +def rename_arguments( + func: Callable[P, R], *, mapper: Mapping[str, str] +) -> Callable[..., R]: ... + + +@overload +def rename_arguments( + *, mapper: Mapping[str, str] +) -> Callable[[Callable[P, R]], Callable[..., R]]: ... + + +def rename_arguments( + func: Callable[P, R] | None = None, + *, + mapper: Mapping[str, str] | None = None, +) -> Callable[..., R] | Callable[[Callable[P, R]], Callable[..., R]]: + """`dags.rename_arguments` with `forwarder=True` baked in. + + See `dags.signature.rename_arguments` for argument semantics. + """ + return _rename_arguments( + func, # ty: ignore[invalid-argument-type] + mapper=mapper, # ty: ignore[invalid-argument-type] + forwarder=True, + ) diff --git a/tests/test_next_state.py b/tests/test_next_state.py index f89bad620..08f51649b 100644 --- a/tests/test_next_state.py +++ b/tests/test_next_state.py @@ -102,7 +102,7 @@ class MockCategory: def test_create_stochastic_next_func(): - labels = jnp.arange(2) + labels = jnp.arange(2, dtype=jnp.int32) got_func = _create_discrete_stochastic_next_func( target="t", next_state_name="next_a", labels=labels ) From 37eb6e6798249605980c8866310807d7821ec778 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Thu, 14 May 2026 09:24:43 +0200 Subject: [PATCH 2/4] Drop _dags_forwarders shim; import dags wrappers directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dags#82 made the `*args, **kwargs` forwarder shape the only behaviour for `with_signature` / `rename_arguments` — there is no `forwarder` flag left to default, so the shim has nothing to do. regime_building imports the dags wrappers directly again; the dags pin moves to the current branch head. Co-Authored-By: Claude Opus 4.7 --- pixi.lock | 46 ++++++------ pyproject.toml | 9 ++- src/lcm/regime_building/Q_and_F.py | 3 +- src/lcm/regime_building/V.py | 3 +- src/lcm/regime_building/max_Q_over_a.py | 2 +- src/lcm/regime_building/next_state.py | 3 +- src/lcm/regime_building/processing.py | 4 +- src/lcm/utils/_dags_forwarders.py | 98 ------------------------- 8 files changed, 34 insertions(+), 134 deletions(-) delete mode 100644 src/lcm/utils/_dags_forwarders.py diff --git a/pixi.lock b/pixi.lock index 3b745fb11..f46905259 100644 --- a/pixi.lock +++ b/pixi.lock @@ -275,7 +275,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2c/1a/aff8bb287a4b1400f69e09a53bd65de96aa5cee5691925b38731c67fc695/click_default_group-1.2.4-py2.py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/8d/2d/f61c918d9edc2127068f0d5ad4604fedd9bfd393f464219090f3279c73f7/estimagic-0.5.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d5/1f/5f4a3cd9e4440e9d9bc78ad0a91a1c8d46b4d429d5239ebe6793c9fe5c41/fsspec-2026.3.0-py3-none-any.whl @@ -588,7 +588,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/78/a3d9ceda0793f4fb43daa292af7b801932611a1aed442636ddfc93d58c7a/jax_cuda12_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -888,7 +888,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/21/98/77f15d81fd0637da454e453c8456d4a2b5c8b2e66823b4237ee8689152cf/jax_cuda13_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -1158,7 +1158,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -1399,7 +1399,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -1637,7 +1637,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-1.3.2-hfd05255_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -1895,7 +1895,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -2146,7 +2146,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -2393,7 +2393,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-1.3.2-hfd05255_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -2648,7 +2648,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl @@ -2911,7 +2911,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -3166,7 +3166,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -3417,7 +3417,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-1.3.2-hfd05255_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -3716,7 +3716,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/78/a3d9ceda0793f4fb43daa292af7b801932611a1aed442636ddfc93d58c7a/jax_cuda12_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -4029,7 +4029,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/21/98/77f15d81fd0637da454e453c8456d4a2b5c8b2e66823b4237ee8689152cf/jax_cuda13_pjrt-0.10.0-py3-none-manylinux_2_27_x86_64.whl @@ -4308,7 +4308,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl @@ -4673,7 +4673,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-ng-2.3.3-hceb46e0_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a1/8e/b2a08ffc51c93842de71f7f988865cebfa7f43d6721957812dc8cc8b9d40/jaxlib-0.10.0-cp314-cp314-manylinux_2_27_x86_64.whl @@ -4964,7 +4964,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zlib-ng-2.3.3-hed4e4f5_1.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/zstd-1.5.7-hbf9d68e_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/a7/25/e1e52a21786b321fb6a2edf9ef9971aa70f06bb2738aef9afd6d8f46a441/jaxlib-0.10.0-cp314-cp314-macosx_11_0_arm64.whl @@ -5274,7 +5274,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/zlib-ng-2.3.3-h0261ad2_1.conda - conda: https://conda.anaconda.org/conda-forge/win-64/zstd-1.5.7-h534d264_6.conda - pypi: https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl - - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 + - pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/70/aa/dfac6d72cc35bc07e7587115b6946e333ef4ccb2e6cd26ecf639438c5d26/jax-0.10.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/24/08/26e6a3ecf0a95f1ec0dcd7a668d5c9a72e581c40fe4ae51e102ca63174c5/jaxlib-0.10.0-cp314-cp314-win_amd64.whl @@ -7542,9 +7542,9 @@ packages: purls: [] size: 210103 timestamp: 1771943128249 -- pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=d415d10#d415d10ee9704b448db90630719f1410ca321132 +- pypi: git+https://github.com/OpenSourceEconomics/dags.git?rev=cf59c04#cf59c04c6ba07b7c54ca763dc155deea3341a480 name: dags - version: 0.5.2.dev1+gd415d10ee + version: 0.5.2.dev6+gcf59c04c6 requires_dist: - flatten-dict - networkx>=3.6 @@ -14085,8 +14085,8 @@ packages: timestamp: 1774796815820 - pypi: ./ name: pylcm - version: 0.0.2.dev179+g59aca3f5f.d20260513 - sha256: 0c5af9810a5be3f1e5861656eaf8e0c6fa09dcbf505d6e746feffeff83e0cce6 + version: 0.0.2.dev182+g1a92dffec.d20260514 + sha256: 2ac3c0e6987658df12a93886cc7ea80b12456826f61aa1720a0c56e8fd828128 requires_dist: - beartype>=0.21 - cloudpickle>=3.1.2 diff --git a/pyproject.toml b/pyproject.toml index e23e07d64..6e8f346d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -150,10 +150,11 @@ jax = ">=0.9" pdbp = "*" pylcm = { path = ".", editable = true } # Pin dags to the feat/no-type-check-flag branch (PR -# OpenSourceEconomics/dags#82): adds the `forwarder=True` flag we need -# to make `with_signature` cooperate with beartype's import claw. -# Will be replaced with a released version once that PR lands. -dags = { git = "https://github.com/OpenSourceEconomics/dags.git", rev = "d415d10" } +# OpenSourceEconomics/dags#82): its wrappers advertise the `*args, +# **kwargs` forwarder shape on `__annotations__`, so beartype's import +# claw treats them as permissive forwarders. Replace with `dags>=0.6` +# once that PR is released. +dags = { git = "https://github.com/OpenSourceEconomics/dags.git", rev = "cf59c04" } [tool.pixi.tasks] asv-compare = "asv compare" asv-preview = "asv preview" diff --git a/src/lcm/regime_building/Q_and_F.py b/src/lcm/regime_building/Q_and_F.py index 9bff01a85..814619ddf 100644 --- a/src/lcm/regime_building/Q_and_F.py +++ b/src/lcm/regime_building/Q_and_F.py @@ -3,7 +3,7 @@ from typing import Any, cast import jax.numpy as jnp -from dags import concatenate_functions +from dags import concatenate_functions, with_signature from jax import Array from lcm.regime_building.h_dag import _get_build_H_kwargs @@ -24,7 +24,6 @@ TransitionFunctionName, TransitionFunctionsMapping, ) -from lcm.utils._dags_forwarders import with_signature from lcm.utils.dispatchers import productmap from lcm.utils.functools import get_union_of_args diff --git a/src/lcm/regime_building/V.py b/src/lcm/regime_building/V.py index ad8c7d21a..31521a5c8 100644 --- a/src/lcm/regime_building/V.py +++ b/src/lcm/regime_building/V.py @@ -3,7 +3,7 @@ from types import MappingProxyType import jax.numpy as jnp -from dags import concatenate_functions +from dags import concatenate_functions, with_signature from dags.tree import qname_from_tree_path from jax import Array @@ -13,7 +13,6 @@ from lcm.regime_building.ndimage import map_coordinates from lcm.shocks import _ShockGrid from lcm.typing import FloatND, ScalarFloat, StateName -from lcm.utils._dags_forwarders import with_signature from lcm.utils.functools import all_as_kwargs from lcm.variables import Variables, get_grids diff --git a/src/lcm/regime_building/max_Q_over_a.py b/src/lcm/regime_building/max_Q_over_a.py index d5fa18b93..f1e5d3913 100644 --- a/src/lcm/regime_building/max_Q_over_a.py +++ b/src/lcm/regime_building/max_Q_over_a.py @@ -3,6 +3,7 @@ from types import MappingProxyType import jax.numpy as jnp +from dags import with_signature from jax import Array from lcm.regime_building.argmax import argmax_and_max @@ -16,7 +17,6 @@ RegimeName, StateName, ) -from lcm.utils._dags_forwarders import with_signature from lcm.utils.dispatchers import productmap diff --git a/src/lcm/regime_building/next_state.py b/src/lcm/regime_building/next_state.py index 6f554c557..5b260ec96 100644 --- a/src/lcm/regime_building/next_state.py +++ b/src/lcm/regime_building/next_state.py @@ -4,7 +4,7 @@ from types import MappingProxyType import jax -from dags import concatenate_functions +from dags import concatenate_functions, with_signature from dags.tree import qname_from_tree_path from jax import Array @@ -27,7 +27,6 @@ TransitionFunctionName, TransitionFunctionsMapping, ) -from lcm.utils._dags_forwarders import with_signature from lcm.variables import Variables diff --git a/src/lcm/regime_building/processing.py b/src/lcm/regime_building/processing.py index 33bb47fd5..7e07e8212 100644 --- a/src/lcm/regime_building/processing.py +++ b/src/lcm/regime_building/processing.py @@ -6,7 +6,8 @@ from typing import Any, Literal, cast import jax -from dags import concatenate_functions, get_annotations +from dags import concatenate_functions, get_annotations, with_signature +from dags.signature import rename_arguments from dags.tree import QNAME_DELIMITER, qname_from_tree_path, tree_path_from_qname from jax import Array from jax import numpy as jnp @@ -64,7 +65,6 @@ UserFunction, VmappedRegimeTransitionFunction, ) -from lcm.utils._dags_forwarders import rename_arguments, with_signature from lcm.utils.containers import ensure_containers_are_immutable from lcm.utils.dispatchers import simulation_spacemap, vmap_1d from lcm.utils.namespace import flatten_regime_namespace, unflatten_regime_namespace diff --git a/src/lcm/utils/_dags_forwarders.py b/src/lcm/utils/_dags_forwarders.py deleted file mode 100644 index 7ba1ab0cb..000000000 --- a/src/lcm/utils/_dags_forwarders.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Thin wrappers around `dags.signature` that default to `forwarder=True`. - -Every direct caller of `dags.with_signature` / `dags.rename_arguments` -inside pylcm produces a generic `*args, **kwargs` forwarder whose -annotations describe the inner function's contract, not the wrapper's -own runtime call protocol. Setting `forwarder=True` on each call tells -beartype's import claw to treat the wrapper as permissive — matching -the wrapper's actual behaviour — instead of enforcing user annotations -against JAX tracers and other forwarded objects. - -Importing these wrappers (rather than the dags ones) keeps that intent -visible at each call site and concentrates the configuration in one -place if pylcm ever needs to flip the default. - -""" - -import inspect -from collections.abc import Callable, Mapping, Sequence -from typing import Any, overload - -from dags import rename_arguments as _rename_arguments -from dags import with_signature as _with_signature -from dags.typing import P, R - - -@overload -def with_signature( - func: Callable[P, R], - *, - args: Mapping[str, str] | Sequence[str] | None = None, - kwargs: Mapping[str, str] | Sequence[str] | None = None, - enforce: bool = True, - return_annotation: Any = inspect.Parameter.empty, # noqa: ANN401 -) -> Callable[P, R]: ... - - -@overload -def with_signature( - *, - args: Mapping[str, str] | Sequence[str] | None = None, - kwargs: Mapping[str, str] | Sequence[str] | None = None, - enforce: bool = True, - return_annotation: Any = inspect.Parameter.empty, # noqa: ANN401 -) -> Callable[[Callable[P, R]], Callable[P, R]]: ... - - -def with_signature( - func: Callable[P, R] | None = None, - *, - args: Mapping[str, str] | Sequence[str] | None = None, - kwargs: Mapping[str, str] | Sequence[str] | None = None, - enforce: bool = True, - return_annotation: Any = inspect.Parameter.empty, -) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]: - """`dags.with_signature` with `forwarder=True` baked in. - - See `dags.signature.with_signature` for argument semantics. The - `forwarder=True` setting advertises the resulting wrapper as a - permissive `*args, **kwargs` forwarder on `__annotations__`, - keeping beartype's claw from enforcing per-parameter annotations - against the forwarder's actual call arguments. - """ - return _with_signature( - func, # ty: ignore[invalid-argument-type] - args=args, - kwargs=kwargs, - enforce=enforce, - return_annotation=return_annotation, - forwarder=True, - ) - - -@overload -def rename_arguments( - func: Callable[P, R], *, mapper: Mapping[str, str] -) -> Callable[..., R]: ... - - -@overload -def rename_arguments( - *, mapper: Mapping[str, str] -) -> Callable[[Callable[P, R]], Callable[..., R]]: ... - - -def rename_arguments( - func: Callable[P, R] | None = None, - *, - mapper: Mapping[str, str] | None = None, -) -> Callable[..., R] | Callable[[Callable[P, R]], Callable[..., R]]: - """`dags.rename_arguments` with `forwarder=True` baked in. - - See `dags.signature.rename_arguments` for argument semantics. - """ - return _rename_arguments( - func, # ty: ignore[invalid-argument-type] - mapper=mapper, # ty: ignore[invalid-argument-type] - forwarder=True, - ) From e9deac348ebb23611405501f16bf54929778aece Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Thu, 14 May 2026 10:47:17 +0200 Subject: [PATCH 3/4] Make GPU-peak-mem subprocess parsing robust to stdout noise The subprocess imports lcm, whose beartype claw can emit diagnostics to stdout, so `int(result.stdout.strip())` blew up on a polluted stream. Mark the peak-bytes line and locate it instead of parsing stdout whole. Co-Authored-By: Claude Opus 4.7 --- benchmarks/_gpu_mem.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/benchmarks/_gpu_mem.py b/benchmarks/_gpu_mem.py index 45e328216..591b6ba18 100644 --- a/benchmarks/_gpu_mem.py +++ b/benchmarks/_gpu_mem.py @@ -26,6 +26,11 @@ class MahlerYumGpuPeakMem(GpuPeakMem): # Project root: the directory containing the benchmarks/ package. _PROJECT_ROOT = Path(__file__).resolve().parent.parent +# Marks the peak-memory line on the subprocess's stdout. The subprocess imports +# lcm, whose beartype claw can emit diagnostics to stdout, so the parent locates +# this line instead of parsing stdout wholesale. +_PEAK_MARKER = "__PEAK_BYTES_IN_USE__" + def measure_gpu_peak(bench_module: str, bench_class: str) -> int: """Run a benchmark in a subprocess and return peak GPU bytes. @@ -58,7 +63,15 @@ def measure_gpu_peak(bench_module: str, bench_class: str) -> int: f"stderr: {result.stderr!r}" ) raise RuntimeError(msg) - return int(result.stdout.strip()) + for line in result.stdout.splitlines(): + if line.startswith(_PEAK_MARKER): + return int(line.removeprefix(_PEAK_MARKER).strip()) + msg = ( + "GPU memory subprocess produced no peak-bytes line.\n" + f"stdout: {result.stdout!r}\n" + f"stderr: {result.stderr!r}" + ) + raise RuntimeError(msg) def _track_gpu_peak_mem(self): @@ -104,4 +117,4 @@ def setup(self): import jax stats = jax.local_devices()[0].memory_stats() - print(stats["peak_bytes_in_use"]) + print(f"{_PEAK_MARKER} {stats['peak_bytes_in_use']}") From bee52c8a39296bb7fea4d826feee745e58752f57 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Thu, 14 May 2026 12:05:47 +0200 Subject: [PATCH 4/4] Make jaxtyping "..." sentinel survive pickling; fix int64 ndimage unit tests Activating the claw on `lcm.regime_building` makes jaxtyping shape checks run on cloudpickled annotation types. jaxtyping marks a `"..."` axis with a plain `object()` sentinel that loses identity across a pickle round-trip, tripping `assert type(variadic_dim) is _NamedVariadicDim`. Replace it with a `__reduce__`-backed singleton, patched in before any jaxtyping-subscripted type is created. Fix `test_ndimage_unit` cases that passed raw int64 arrays the int32-pinned primitives now reject. Co-Authored-By: Claude Opus 4.7 --- src/lcm/__init__.py | 4 ++++ src/lcm/_jaxtyping_patch.py | 36 ++++++++++++++++++++++++++++++++++++ tests/test_ndimage_unit.py | 13 ++++++++----- 3 files changed, 48 insertions(+), 5 deletions(-) create mode 100644 src/lcm/_jaxtyping_patch.py diff --git a/src/lcm/__init__.py b/src/lcm/__init__.py index 565acfce4..e18fdc350 100644 --- a/src/lcm/__init__.py +++ b/src/lcm/__init__.py @@ -25,6 +25,10 @@ import jax +# Patch jaxtyping's `"..."` sentinel to survive pickling before any +# `jaxtyping`-subscripted type is created (see the module docstring). +from lcm import _jaxtyping_patch # noqa: F401 + with contextlib.suppress(ImportError): import pdbp # noqa: F401 diff --git a/src/lcm/_jaxtyping_patch.py b/src/lcm/_jaxtyping_patch.py new file mode 100644 index 000000000..e9dced877 --- /dev/null +++ b/src/lcm/_jaxtyping_patch.py @@ -0,0 +1,36 @@ +"""Make jaxtyping's anonymous-variadic-dim sentinel survive pickling. + +jaxtyping marks a `"..."` axis with a module-level `object()` sentinel +(`_anonymous_variadic_dim`). A plain `object()` does not keep its identity +across a pickle round-trip, so cloudpickling a value whose type annotations +reference a `Foo[Array, "..."]` type — which the beartype claw makes +pervasive — yields a type whose variadic-dim marker no longer matches the +live module global. jaxtyping's shape check then trips +`assert type(variadic_dim) is _NamedVariadicDim`. + +Replacing the sentinel with a `__reduce__`-backed singleton makes it +round-trip to the same object, so unpickled annotation types stay valid. +This module must be imported before any `jaxtyping`-subscripted type is +created — `lcm/__init__.py` imports it before every other `lcm` submodule. +""" + +from typing import Self + +from jaxtyping import _array_types + + +class _AnonymousVariadicDim: + """Picklable singleton for jaxtyping's `"..."` axis marker.""" + + _instance: Self | None = None + + def __new__(cls) -> Self: + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __reduce__(self) -> tuple[type[_AnonymousVariadicDim], tuple[()]]: + return (_AnonymousVariadicDim, ()) + + +_array_types._anonymous_variadic_dim = _AnonymousVariadicDim() # noqa: SLF001 diff --git a/tests/test_ndimage_unit.py b/tests/test_ndimage_unit.py index 516a5aec2..f70903acc 100644 --- a/tests/test_ndimage_unit.py +++ b/tests/test_ndimage_unit.py @@ -12,8 +12,11 @@ def test_map_coordinates_wrong_input_dimensions(): - values = jnp.arange(2) # ndim = 1 - coordinates = [jnp.array([0]), jnp.array([1])] # len = 2 + values = jnp.arange(2, dtype=jnp.int32) # ndim = 1 + coordinates = [ + jnp.array([0], dtype=jnp.int32), + jnp.array([1], dtype=jnp.int32), + ] # len = 2 with pytest.raises(ValueError, match="coordinates must be a sequence of length"): map_coordinates(values, coordinates) @@ -29,7 +32,7 @@ def test_map_coordinates_extrapolation(): def test_nonempty_sum(): - a = jnp.arange(3) + a = jnp.arange(3, dtype=jnp.int32) expected = a + a + a got = _sum_all([a, a, a]) @@ -38,7 +41,7 @@ def test_nonempty_sum(): def test_nonempty_prod(): - a = jnp.arange(3) + a = jnp.arange(3, dtype=jnp.int32) expected = a * a * a got = _multiply_all([a, a, a]) @@ -75,7 +78,7 @@ def test_linear_indices_and_weights_inside_domain(): def test_linear_indices_and_weights_outside_domain(): - coordinates = jnp.array([-1, 2]) + coordinates = jnp.array([-1.0, 2.0]) (idx_low, weight_low), (idx_high, weight_high) = _compute_indices_and_weights( coordinates, input_size=2