test python 3.14 (#2157)

Author: juanitorduz

.github/workflows/ci.yml

@@ -25,7 +25,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.11", "3.13"]
+        python-version: ["3.11", "3.14"]
 
     steps:
       - uses: actions/checkout@v6
@@ -60,7 +60,7 @@ jobs:
     needs: [lint, prek]
     strategy:
       matrix:
-        python-version: ["3.11", "3.13"]
+        python-version: ["3.11", "3.14"]
     env:
       UV_PYTHON: ${{ matrix.python-version }}
 
@@ -88,7 +88,7 @@ jobs:
         run: |
           JAX_ENABLE_X64=1 uv run pytest -vs test/test_distributions.py -k "powerLaw or Dagum"
       - name: Test tracer leak
-        if: matrix.python-version == '3.13'
+        if: matrix.python-version == '3.14'
         env:
           JAX_CHECK_TRACER_LEAKS: 1
         run: |
@@ -99,7 +99,7 @@ jobs:
             test/infer/test_mcmc.py::test_reuse_mcmc_run
           uv run pytest -vs test/test_distributions.py::test_mean_var -k Gompertz
       - name: Coveralls
-        if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.13'
+        if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.14'
         uses: coverallsapp/github-action@v2
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -112,7 +112,7 @@ jobs:
     needs: [lint, prek]
     strategy:
       matrix:
-        python-version: ["3.11", "3.13"]
+        python-version: ["3.11", "3.14"]
     env:
       UV_PYTHON: ${{ matrix.python-version }}
 
@@ -155,7 +155,7 @@ jobs:
         run: |
           JAX_ENABLE_X64=1 uv run pytest -vs test/contrib/test_nested_sampling.py
       - name: Coveralls
-        if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.13'
+        if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.14'
         uses: coverallsapp/github-action@v2
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
@@ -168,7 +168,7 @@ jobs:
     needs: [lint, prek]
     strategy:
       matrix:
-        python-version: ["3.13"]
+        python-version: ["3.14"]
     env:
       UV_PYTHON: ${{ matrix.python-version }}
 
@@ -194,7 +194,7 @@ jobs:
         run: |
           CI=1 XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs -k test_example
       - name: Coveralls
-        if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.13'
+        if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.14'
         uses: coverallsapp/github-action@v2
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}

numpyro/contrib/nested_sampling.py

@@ -21,7 +21,7 @@
     from jaxns.public import DefaultNestedSampler
     from jaxns.utils import NestedSamplerResults
 
-except ImportError as e:
+except (ImportError, AttributeError) as e:
     raise ImportError(
         f"{e} \n "
         f"To use this module, please install `jaxns>2.5` package. It can be"

numpyro/examples/datasets.py

@@ -172,15 +172,21 @@ def _download_with_retries(url: str, out_path: str) -> None:
             request = Request(url, headers={"User-Agent": "numpyro-datasets"})
             with urlopen(request) as response, open(out_path, "wb") as f:
                 shutil.copyfileobj(response, f)
+            if isinstance(last_exc, HTTPError):
+                last_exc.close()
             return
         except (HTTPError, URLError) as exc:
             retryable = _is_retryable_error(exc)
             delay = _download_backoff_delay(exc, attempt) if retryable else None
             detached_exc = _detached_download_error(exc)
             if isinstance(exc, HTTPError):
                 exc.close()
+            if isinstance(last_exc, HTTPError):
+                last_exc.close()
             last_exc = detached_exc
             if not retryable:
+                if isinstance(detached_exc, HTTPError):
+                    detached_exc.close()
                 raise detached_exc
             if attempt < _DOWNLOAD_MAX_RETRIES - 1:
                 print(
@@ -189,6 +195,8 @@ def _download_with_retries(url: str, out_path: str) -> None:
                     )
                 )
                 time.sleep(delay)
+    if isinstance(last_exc, HTTPError):
+        last_exc.close()
     raise last_exc
 
 
@@ -221,6 +229,8 @@ def _download(dset: dset) -> None:
             print("Download complete.")
             break
         else:
+            if isinstance(last_exc, HTTPError):
+                last_exc.close()
             raise last_exc
 
 

test/contrib/test_nested_sampling.py

@@ -11,22 +11,23 @@
 import jax.numpy as jnp
 
 import numpyro
+import numpyro.distributions as dist
+from numpyro.distributions.transforms import AffineTransform, ExpTransform
 
+_jaxns_available = True
 try:
     if os.environ.get("JAX_ENABLE_X64"):
         from numpyro.contrib.nested_sampling import NestedSampler, UniformReparam
 
 except ImportError:
-    pytestmark = pytest.mark.skip(reason="jaxns is not installed")
+    _jaxns_available = False
 
-import numpyro.distributions as dist
-from numpyro.distributions.transforms import AffineTransform, ExpTransform
 
 pytestmark = [
     pytest.mark.filterwarnings("ignore:jax.tree_.+ is deprecated:FutureWarning"),
     pytest.mark.filterwarnings("ignore:JAX x64"),
     pytest.mark.skipif(
-        not os.environ.get("JAX_ENABLE_X64"),
+        not os.environ.get("JAX_ENABLE_X64") or not _jaxns_available,
         reason="test suite for jaxns requires double precision",
     ),
 ]