From bd2395401a141f1f8c57ebed4bd9674b0fa01f0b Mon Sep 17 00:00:00 2001 From: beckermr Date: Fri, 17 Jan 2025 13:16:49 -0600 Subject: [PATCH 1/5] feat: add benchmark for spergel init --- tests/jax/test_benchmarks.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/jax/test_benchmarks.py b/tests/jax/test_benchmarks.py index 7ef21332..76b45a41 100644 --- a/tests/jax/test_benchmarks.py +++ b/tests/jax/test_benchmarks.py @@ -212,3 +212,16 @@ def test_benchmark_spergel_kvalue(benchmark, kind): benchmark, kind, lambda: _run_spergel_bench_kvalue_jit().block_until_ready() ) print(f"time: {dt:0.4g} ms", end=" ") + + +@jax.jit +def _run_spergel_bench_init(): + return jgs.Spergel(nu=-0.6, half_light_radius=3.4).scale_radius + + +@pytest.mark.parametrize("kind", ["compile", "run"]) +def test_benchmark_spergel_init(benchmark, kind): + dt = _run_benchmarks( + benchmark, kind, lambda: _run_spergel_bench_init().block_until_ready() + ) + print(f"time: {dt:0.4g} ms", end=" ") From dfab96064fe8615d919ecb6eb91469ee85ec1869 Mon Sep 17 00:00:00 2001 From: beckermr Date: Fri, 17 Jan 2025 13:18:15 -0600 Subject: [PATCH 2/5] feat: add gaussian init too --- tests/jax/test_benchmarks.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/jax/test_benchmarks.py b/tests/jax/test_benchmarks.py index 76b45a41..d00cad23 100644 --- a/tests/jax/test_benchmarks.py +++ b/tests/jax/test_benchmarks.py @@ -225,3 +225,16 @@ def test_benchmark_spergel_init(benchmark, kind): benchmark, kind, lambda: _run_spergel_bench_init().block_until_ready() ) print(f"time: {dt:0.4g} ms", end=" ") + + +@jax.jit +def _run_gaussian_bench_init(): + return jgs.Gaussian(half_light_radius=3.4).sigma + + +@pytest.mark.parametrize("kind", ["compile", "run"]) +def test_benchmark_gaussian_init(benchmark, kind): + dt = _run_benchmarks( + benchmark, kind, lambda: _run_gaussian_bench_init().block_until_ready() + ) + print(f"time: {dt:0.4g} ms", end=" ") From ac0cfeccc4dfa68dee62ed14b89d10d667aab2f8 Mon Sep 17 00:00:00 2001 From: beckermr Date: Fri, 17 Jan 2025 15:23:07 -0600 Subject: [PATCH 3/5] fix: pin back to 0.4.* for now --- .github/workflows/benchmarks.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/benchmarks.yaml b/.github/workflows/benchmarks.yaml index 12b1244d..07394336 100644 --- a/.github/workflows/benchmarks.yaml +++ b/.github/workflows/benchmarks.yaml @@ -35,6 +35,8 @@ jobs: python -m pip install --upgrade pip python -m pip install pytest pytest-codspeed python -m pip install . + # temp pin until 0.5 is on conda + python -m pip install jax~=0.4 git submodule update --init --recursive - name: Run benchmarks From d00fdcc9ea925a8b4153fdeea088e2b5998846e6 Mon Sep 17 00:00:00 2001 From: beckermr Date: Fri, 17 Jan 2025 15:23:54 -0600 Subject: [PATCH 4/5] fix: wrong workflow file --- .github/workflows/benchmarks.yaml | 2 -- .github/workflows/python_package.yaml | 4 +++- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/benchmarks.yaml b/.github/workflows/benchmarks.yaml index 07394336..12b1244d 100644 --- a/.github/workflows/benchmarks.yaml +++ b/.github/workflows/benchmarks.yaml @@ -35,8 +35,6 @@ jobs: python -m pip install --upgrade pip python -m pip install pytest pytest-codspeed python -m pip install . - # temp pin until 0.5 is on conda - python -m pip install jax~=0.4 git submodule update --init --recursive - name: Run benchmarks diff --git a/.github/workflows/python_package.yaml b/.github/workflows/python_package.yaml index 399fdbd8..5919ae43 100644 --- a/.github/workflows/python_package.yaml +++ b/.github/workflows/python_package.yaml @@ -20,7 +20,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.9", "3.10", "3.11", "3.12"] + python-version: ["3.10", "3.11", "3.12"] steps: - uses: actions/checkout@v4 @@ -35,6 +35,8 @@ jobs: python -m pip install --upgrade pip python -m pip install pytest pytest-codspeed python -m pip install . + # temp pin until 0.5 is on conda + python -m pip install jax~=0.4 - name: Test with pytest run: | From f56e1abaccffc51470f80c90249f6b36a20f1e37 Mon Sep 17 00:00:00 2001 From: beckermr Date: Fri, 17 Jan 2025 15:37:13 -0600 Subject: [PATCH 5/5] fix: idk what pip[ is smoking 0.5.0 is not compatible ~=0.4 but w/e try this --- .github/workflows/python_package.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/python_package.yaml b/.github/workflows/python_package.yaml index 5919ae43..7c0d30ff 100644 --- a/.github/workflows/python_package.yaml +++ b/.github/workflows/python_package.yaml @@ -36,7 +36,7 @@ jobs: python -m pip install pytest pytest-codspeed python -m pip install . # temp pin until 0.5 is on conda - python -m pip install jax~=0.4 + python -m pip install "jax<0.5.0" - name: Test with pytest run: |