From f55baf35e4f314c49b17cd61c0b475f743fb39a8 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 2 Sep 2025 13:01:57 +0000 Subject: [PATCH 01/16] Bump actions/checkout from 4 to 5 Bumps [actions/checkout](https://github.com/actions/checkout) from 4 to 5. - [Release notes](https://github.com/actions/checkout/releases) - [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md) - [Commits](https://github.com/actions/checkout/compare/v4...v5) --- updated-dependencies: - dependency-name: actions/checkout dependency-version: '5' dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- .github/workflows/ci.yml | 4 ++-- .github/workflows/pypi.yml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 04960d8e..015c028d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,7 +25,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 with: # setuptools_scm requires a non-shallow clone of the repository fetch-depth: 0 @@ -52,7 +52,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index 3bc4b12b..9c4640bf 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -14,7 +14,7 @@ jobs: id-token: write steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 with: # setuptools_scm requires a non-shallow clone of the repository fetch-depth: 0 From 31934cc9dafa2a5dda46c708db7c253b0c46edb5 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 1 Oct 2025 21:05:42 +0000 Subject: [PATCH 02/16] Bump pypa/gh-action-pypi-publish from 1.12.4 to 1.13.0 Bumps [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) from 1.12.4 to 1.13.0. - [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases) - [Commits](https://github.com/pypa/gh-action-pypi-publish/compare/v1.12.4...v1.13.0) --- updated-dependencies: - dependency-name: pypa/gh-action-pypi-publish dependency-version: 1.13.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- .github/workflows/pypi.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index 3bc4b12b..60b47127 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -25,4 +25,4 @@ jobs: - name: Build SDist run: pipx run build --sdist - - uses: pypa/gh-action-pypi-publish@v1.12.4 + - uses: pypa/gh-action-pypi-publish@v1.13.0 From fe864a517ed0d4f68b02a5e1f6c7662bdf9828eb Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 16 Oct 2025 15:02:23 +0000 Subject: [PATCH 03/16] Bump actions/setup-python from 5 to 6 Bumps [actions/setup-python](https://github.com/actions/setup-python) from 5 to 6. - [Release notes](https://github.com/actions/setup-python/releases) - [Commits](https://github.com/actions/setup-python/compare/v5...v6) --- updated-dependencies: - dependency-name: actions/setup-python dependency-version: '6' dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] --- .github/workflows/ci.yml | 4 ++-- .github/workflows/pypi.yml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 015c028d..86ebc55f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,7 +31,7 @@ jobs: fetch-depth: 0 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} @@ -55,7 +55,7 @@ jobs: uses: actions/checkout@v5 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: "3.11" diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index 9c4640bf..094810d6 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -19,7 +19,7 @@ jobs: # setuptools_scm requires a non-shallow clone of the repository fetch-depth: 0 - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v6 name: Install Python - name: Build SDist From 3ffaa3502aa38e138838f61a7ffcba18faabb3ae Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:00:57 +0200 Subject: [PATCH 04/16] Exclude jax[cpu] version 0.8.0 from dependencies --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 24a63268..bf644374 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.8.0", "interpax", "astroquery", "beartype", From 670b9ea51a981bdfe34438a884a942a6909535fb Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:04:45 +0200 Subject: [PATCH 05/16] Exclude jax[cpu] version 0.7.2 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index bf644374..706d94fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.7.2", "jax[cpu]!=0.8.0", "interpax", "astroquery", From ce595d35016ce074681d437f48ef19227fc10319 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:12:59 +0200 Subject: [PATCH 06/16] Update pyproject.toml --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 706d94fc..00a5830e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.7.1", "jax[cpu]!=0.7.2", "jax[cpu]!=0.8.0", "interpax", From 9d908a51683d4323caf4213b88edd1438aef88bd Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:16:00 +0200 Subject: [PATCH 07/16] Exclude jax[cpu] version 0.7.0 Add exclusion for jax[cpu] version 0.7.0 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 00a5830e..abe7c96b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.7.0", "jax[cpu]!=0.7.1", "jax[cpu]!=0.7.2", "jax[cpu]!=0.8.0", From ae60edf3ab7aff554b1f817d7ed612da6fefbcb2 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:19:41 +0200 Subject: [PATCH 08/16] Exclude jax[cpu] version 0.6.2 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index abe7c96b..442b1af1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.6.2", "jax[cpu]!=0.7.0", "jax[cpu]!=0.7.1", "jax[cpu]!=0.7.2", From 71001b0b2c61336c4a79d84bd3157dcb6da20839 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:22:55 +0200 Subject: [PATCH 09/16] Exclude jax[cpu] version 0.6.1 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 442b1af1..c319a5db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.6.1", "jax[cpu]!=0.6.2", "jax[cpu]!=0.7.0", "jax[cpu]!=0.7.1", From b35c48e47166c2d02e34e0e7a78e95b96ebf4815 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:27:32 +0200 Subject: [PATCH 10/16] Exclude jax[cpu] version 0.6.0 from dependencies --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index c319a5db..68e036ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.6.0", "jax[cpu]!=0.6.1", "jax[cpu]!=0.6.2", "jax[cpu]!=0.7.0", From 17a4023baef43e279353100fec541e1cb07526e1 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:31:05 +0200 Subject: [PATCH 11/16] Exclude jax[cpu] version 0.5.3 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 68e036ad..ea6b7015 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.5.3", "jax[cpu]!=0.6.0", "jax[cpu]!=0.6.1", "jax[cpu]!=0.6.2", From 25412a4a962f0cf3a9ace235d7bcda34a357a432 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 00:36:31 +0200 Subject: [PATCH 12/16] Exclude jax[cpu] version 0.5.2 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index ea6b7015..07aaf294 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", + "jax[cpu]!=0.5.2", "jax[cpu]!=0.5.3", "jax[cpu]!=0.6.0", "jax[cpu]!=0.6.1", From 6ea1eccf8e1f44189c00ef43ffc63c37018357e7 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 10:05:21 +0200 Subject: [PATCH 13/16] Clean up jax[cpu] version exclusions in pyproject.toml Removed several version exclusions for jax[cpu] in dependencies. --- pyproject.toml | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 07aaf294..da54c1b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,14 +35,6 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", - "jax[cpu]!=0.5.2", - "jax[cpu]!=0.5.3", - "jax[cpu]!=0.6.0", - "jax[cpu]!=0.6.1", - "jax[cpu]!=0.6.2", - "jax[cpu]!=0.7.0", - "jax[cpu]!=0.7.1", - "jax[cpu]!=0.7.2", "jax[cpu]!=0.8.0", "interpax", "astroquery", @@ -62,6 +54,7 @@ tests = [ "nbval", "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", + "jax[cpu]!=0.8.0", "pre-commit", ] docs = [ From 39db4a3d3d195b8642f811d46703d05a01d2821f Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 10:59:38 +0200 Subject: [PATCH 14/16] Add typechecker to _cumtrapz_scan_func decorator --- rubix/cosmology/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rubix/cosmology/utils.py b/rubix/cosmology/utils.py index 0579fec8..11791a7f 100644 --- a/rubix/cosmology/utils.py +++ b/rubix/cosmology/utils.py @@ -8,8 +8,8 @@ # Source: https://github.com/ArgonneCPAC/dsps/blob/b81bac59e545e2d68ccf698faba078d87cfa2dd8/dsps/utils.py#L247C1-L256C1 -@jaxtyped(typechecker=typechecker) @jit +@jaxtyped(typechecker=typechecker) def _cumtrapz_scan_func(carryover, el): """ Integral helper function, which uses the formula for trapezoidal integration. From 2fefe489c25dff57d3aad38911a54f40b8a7e7c8 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 20:26:19 +0200 Subject: [PATCH 15/16] switch @jit and typechecker decorator to make tests pass. --- pyproject.toml | 2 -- rubix/cosmology/utils.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index da54c1b6..24a63268 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,6 @@ dependencies = [ "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", "jax[cpu]!=0.5.1", - "jax[cpu]!=0.8.0", "interpax", "astroquery", "beartype", @@ -54,7 +53,6 @@ tests = [ "nbval", "jax[cpu]!=0.4.27", "jax[cpu]!=0.4.36", - "jax[cpu]!=0.8.0", "pre-commit", ] docs = [ diff --git a/rubix/cosmology/utils.py b/rubix/cosmology/utils.py index 11791a7f..60a6f9d9 100644 --- a/rubix/cosmology/utils.py +++ b/rubix/cosmology/utils.py @@ -37,8 +37,8 @@ def _cumtrapz_scan_func(carryover, el): # Source: https://github.com/ArgonneCPAC/dsps/blob/b81bac59e545e2d68ccf698faba078d87cfa2dd8/dsps/utils.py#L278C1-L298C1 -@jaxtyped(typechecker=typechecker) @jit +@jaxtyped(typechecker=typechecker) def trapz( xarr: Union[jnp.ndarray, Float[Array, "n"]], yarr: Union[jnp.ndarray, Float[Array, "n"]], From 81b5671398176aa80843b81aada7e1446ba29448 Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sat, 25 Oct 2025 20:35:17 +0200 Subject: [PATCH 16/16] switch @jit with typechecker decorator to make tests pass. --- rubix/cosmology/base.py | 31 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/rubix/cosmology/base.py b/rubix/cosmology/base.py index b5ce7d24..1790b716 100644 --- a/rubix/cosmology/base.py +++ b/rubix/cosmology/base.py @@ -56,8 +56,8 @@ def __init__(self, Om0: float, w0: float, wa: float, h: float): self.wa = jnp.float32(wa) self.h = jnp.float32(h) - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def scale_factor_to_redshift( self, a: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -79,8 +79,8 @@ def scale_factor_to_redshift( z = 1.0 / a - 1.0 return z - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def _rho_de_z(self, z: Union[Float[Array, "..."], float]) -> Float[Array, "..."]: a = 1.0 / (1.0 + z) de_z = a ** (-3.0 * (1.0 + self.w0 + self.wa)) * lax.exp( @@ -88,8 +88,8 @@ def _rho_de_z(self, z: Union[Float[Array, "..."], float]) -> Float[Array, "..."] ) return de_z - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def _Ez(self, z: Union[Float[Array, "..."], float]) -> Float[Array, "..."]: zp1 = 1.0 + z Ode0 = 1.0 - self.Om0 @@ -97,15 +97,15 @@ def _Ez(self, z: Union[Float[Array, "..."], float]) -> Float[Array, "..."]: E = jnp.sqrt(t) return E - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def _integrand_oneOverEz( self, z: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: return 1 / self._Ez(z) - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def comoving_distance_to_z( self, redshift: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -128,8 +128,8 @@ def comoving_distance_to_z( integrand = self._integrand_oneOverEz(z_table) return trapz(z_table, integrand) * C_SPEED * 1e-5 / self.h - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def luminosity_distance_to_z( self, redshift: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -150,8 +150,8 @@ def luminosity_distance_to_z( """ return self.comoving_distance_to_z(redshift) * (1 + redshift) - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def angular_diameter_distance_to_z( self, redshift: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -172,8 +172,8 @@ def angular_diameter_distance_to_z( """ return self.comoving_distance_to_z(redshift) / (1 + redshift) - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def distance_modulus_to_z( self, redshift: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -196,15 +196,15 @@ def distance_modulus_to_z( mu = 5.0 * jnp.log10(d_lum * 1e5) return mu - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def _hubble_time(self, z: Union[Float[Array, "..."], float]) -> Float[Array, "..."]: E0 = self._Ez(z) htime = 1e-16 * MPC / YEAR / self.h / E0 return htime - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def lookback_to_z( self, redshift: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -229,8 +229,8 @@ def lookback_to_z( th = self._hubble_time(0.0) return th * res - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def age_at_z0(self) -> Float[Array, "..."]: """ The function calculates the age of the universe at redshift 0. @@ -250,8 +250,8 @@ def age_at_z0(self) -> Float[Array, "..."]: th = self._hubble_time(0.0) return th * res - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def _age_at_z_kern( self, redshift: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -259,8 +259,8 @@ def _age_at_z_kern( tlook = self.lookback_to_z(redshift) return t0 - tlook - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def age_at_z( self, redshift: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -285,8 +285,8 @@ def age_at_z( def _age_at_z_vmap(self): return jit(vmap(self._age_at_z_kern)) - @jaxtyped(typechecker=typechecker) @jit + @jaxtyped(typechecker=typechecker) def angular_scale( self, z: Union[Float[Array, "..."], float] ) -> Float[Array, "..."]: @@ -327,9 +327,6 @@ def _Om_at_z(self, z): E = self._Ez(z) return self.Om0 * (1.0 + z) ** 3 / E / E - - - @jit def _delta_vir(self, z): x = self._Om(z) - 1.0