build: refactor build infrastruture

Signed-off-by: nstarman <[email protected]>

.pre-commit-config.yaml

@@ -20,12 +20,12 @@ repos:
       - id: debug-statements
       - id: check-yaml
       - id: check-added-large-files
-        exclude: notebooks/*
+        exclude: notebooks/
 
   - repo: https://github.com/codespell-project/codespell
     rev: v2.3.0
     hooks:
       - id: codespell
-        stages: [commit, commit-msg]
+        stages: [pre-commit, commit-msg]
         args:
           [--ignore-words-list, "Teh,aas", --check-filenames, --skip, "*.ipynb"]

.python-version

@@ -0,0 +1 @@
+3.9

pyproject.toml

@@ -1,24 +1,101 @@
+[project]
+name = "numpyro"
+version = "0.15.3"
+description = "Pyro PPL on NumPy"
+readme = "README.md"
+requires-python = ">=3.9"
+authors = [
+    { name = "Uber AI Labs", email = "[email protected]" },
+]
+keywords = ["probabilistic", "machine learning", "bayesian", "statistics"]
+license.file = "LICENSE.md"
+classifiers=[
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: POSIX :: Linux",
+    "Operating System :: MacOS :: MacOS X",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+]
+dependencies = [
+    "jax>=0.4.25",
+    "jaxlib>=0.4.25",
+    "multipledispatch>=1.0.0",
+    "numpy>=2.0.2",
+    "tqdm>=4.66.5",
+]
+
+[project.optional-dependencies]
+doc = [
+    "ipython",  # sphinx needs this to render codes
+    "nbsphinx>=0.8.9",
+    "readthedocs-sphinx-search>=0.3.2",
+    "sphinx>=5",
+    "sphinx_rtd_theme",
+    "sphinx-gallery",
+]
+test = [
+    "importlib-metadata<5.0",
+    "ruff>=0.1.8",
+    "pytest>=4.1",
+    "pyro-api>=0.1.1",
+    "scikit-learn",
+    "scipy>=1.9",
+]
+dev = [
+    "dm-haiku",
+    "flax",
+    "funsor>=0.4.1",
+    "graphviz",
+    "jaxns==2.4.8",
+    "matplotlib",
+    "optax>=0.0.6",
+    "pylab-sdk",  # jaxns dependency
+    "pyyaml",  # flax dependency
+    "requests",  # pylab dependency
+    "tensorflow_probability>=0.18.0",
+]
+examples = [
+    "arviz",
+    "jupyter",
+    "matplotlib",
+    "pandas",
+    "seaborn",
+    "scikit-learn",
+    "wordcloud",
+]
+cpu = ["numpyro", "jax[cpu]"]
+# gpu = ["numpyro", "jax[tpu]"]
+# cuda = ["numpyro", "jax[cuda]"]
+
+[project.urls]
+Documentation = "https://num.pyro.ai/en/stable/"
+Repository = "https://github.com/pyro-ppl/numpyro"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+include = ["src/numpyro/version.py"]
+
+[tool.hatch.build.targets.sdist.hooks.custom]
+hooks = ["write_version_file"]
+
+[tool.hatch.hooks.write_version_file]
+path = "your_package/version.py"
+template = '''
+# Auto-generated version file
+__version__ = "{version}"
+'''
+
 [tool.ruff]
 # Exclude a variety of commonly ignored directories.
 exclude = [
-    ".bzr",
-    ".direnv",
-    ".eggs",
-    ".git",
-    ".git-rewrite",
-    ".hg",
-    ".ipynb_checkpoints",
-    ".mypy_cache",
-    ".nox",
-    ".pants.d",
-    ".pyenv",
-    ".pytest_cache",
-    ".pytype",
-    ".ruff_cache",
-    ".svn",
-    ".tox",
-    ".venv",
-    ".vscode",
+    ".*",
     "__pypackages__",
     "_build",
     "buck-out",
@@ -63,8 +140,8 @@ skip-magic-trailing-comma = false
 line-ending = "auto"
 
 [tool.ruff.lint.extend-per-file-ignores]
-"numpyro/contrib/tfp/distributions.py" = ["F811"]
-"numpyro/distributions/kl.py" = ["F811"]
+"src/numpyro/contrib/tfp/distributions.py" = ["F811"]
+"src/numpyro/distributions/kl.py" = ["F811"]
 
 [tool.ruff.lint.isort]
 combine-as-imports = true
@@ -102,3 +179,8 @@ doctest_optionflags = [
     "NORMALIZE_WHITESPACE",
     "IGNORE_EXCEPTION_DETAIL",
 ]
+
+[tool.uv]
+dev-dependencies = [
+    "pre-commit>=4.0.1",
+]

numpyro/contrib/control_flow/scan.py → src/numpyro/contrib/control_flow/scan.py

@@ -197,9 +197,10 @@ def body_fn(wrapped_carry, x, prefix=None):
             )
         return (i + 1, rng_key, new_carry), (PytreeTrace(trace), y)
 
-    with handlers.block(
-        hide_fn=lambda site: not site["name"].startswith("_PREV_")
-    ), enum(first_available_dim=first_available_dim):
+    with (
+        handlers.block(hide_fn=lambda site: not site["name"].startswith("_PREV_")),
+        enum(first_available_dim=first_available_dim),
+    ):
         wrapped_carry = (0, rng_key, init)
         y0s = []
         # We run unroll_steps + 1 where the last step is used for rolling with `lax.scan`

numpyro/infer/autoguide.py → src/numpyro/infer/autoguide.py

@@ -1463,8 +1463,10 @@ def _sample_latent(self, *args, **kwargs):
         if self.global_guide is not None:
             global_latents = self.global_guide(*args, **kwargs)
             rng_key = numpyro.prng_key()
-            with handlers.block(), handlers.seed(rng_seed=rng_key), handlers.substitute(
-                data=global_latents
+            with (
+                handlers.block(),
+                handlers.seed(rng_seed=rng_key),
+                handlers.substitute(data=global_latents),
             ):
                 global_outputs = self.global_guide.model(*args, **kwargs)
             local_args = (global_outputs,)
@@ -1575,9 +1577,12 @@ def fn(x):
             if self.local_guide is not None:
                 key = numpyro.prng_key()
                 subsample_guide = partial(_subsample_model, self.local_guide)
-                with handlers.block(), handlers.trace() as tr, handlers.seed(
-                    rng_seed=key
-                ), handlers.substitute(data=local_guide_params):
+                with (
+                    handlers.block(),
+                    handlers.trace() as tr,
+                    handlers.seed(rng_seed=key),
+                    handlers.substitute(data=local_guide_params),
+                ):
                     with warnings.catch_warnings():
                         warnings.simplefilter("ignore")
                         subsample_guide(*local_args, **local_kwargs)

numpyro/infer/elbo.py → src/numpyro/infer/elbo.py

@@ -874,8 +874,11 @@ def get_importance_trace_enum(
         trace as _trace,
     )
 
-    with plate_to_enum_plate(), enum(
-        first_available_dim=(-max_plate_nesting - 1) if max_plate_nesting else None
+    with (
+        plate_to_enum_plate(),
+        enum(
+            first_available_dim=(-max_plate_nesting - 1) if max_plate_nesting else None
+        ),
     ):
         guide = substitute(guide, data=params)
         with _without_rsample_stop_gradient():

numpyro/infer/inspect.py → src/numpyro/infer/inspect.py

@@ -58,8 +58,10 @@ def get_trace():
 
 def _get_log_probs(model, model_args, model_kwargs, **sample):
     # Note: We use seed 0 for parameter initialization.
-    with handlers.trace() as tr, handlers.seed(rng_seed=0), handlers.substitute(
-        data=sample
+    with (
+        handlers.trace() as tr,
+        handlers.seed(rng_seed=0),
+        handlers.substitute(data=sample),
     ):
         model(*model_args, **model_kwargs)
     return {
@@ -370,8 +372,9 @@ def process_message(self, msg):
 
         # Note: We use seed 0 for parameter initialization.
         with handlers.trace() as tr, handlers.seed(rng_seed=0):
-            with handlers.substitute(data=sample), substitute_deterministic(
-                data=sample
+            with (
+                handlers.substitute(data=sample),
+                substitute_deterministic(data=sample),
             ):
                 model(*model_args, **model_kwargs)
         provenance_arrays = {}

test/infer/test_hmc_util.py

@@ -57,8 +57,9 @@ def optimize(f):
 @pytest.mark.parametrize("regularize", [True, False])
 @pytest.mark.filterwarnings("ignore:numpy.linalg support is experimental:UserWarning")
 def test_welford_covariance(jitted, diagonal, regularize):
-    with optional(jitted, disable_jit()), optional(
-        jitted, control_flow_prims_disabled()
+    with (
+        optional(jitted, disable_jit()),
+        optional(jitted, control_flow_prims_disabled()),
     ):
         np.random.seed(0)
         loc = np.random.randn(3)

test/test_distributions.py

@@ -2221,8 +2221,12 @@ def g(x):
 
 
 def test_beta_proportion_invalid_mean():
-    with dist.distribution.validation_enabled(), pytest.raises(
-        ValueError, match=r"^BetaProportion distribution got invalid mean parameter\.$"
+    with (
+        dist.distribution.validation_enabled(),
+        pytest.raises(
+            ValueError,
+            match=r"^BetaProportion distribution got invalid mean parameter\.$",
+        ),
     ):
         dist.BetaProportion(1.0, 1.0)
 

test/test_handlers.py

@@ -372,8 +372,10 @@ def test_subsample_substitute():
     data = jnp.arange(100.0)
     subsample_size = 7
     subsample = jnp.array([13, 3, 30, 4, 1, 68, 5])
-    with handlers.trace() as tr, handlers.seed(rng_seed=0), handlers.substitute(
-        data={"a": subsample}
+    with (
+        handlers.trace() as tr,
+        handlers.seed(rng_seed=0),
+        handlers.substitute(data={"a": subsample}),
     ):
         with numpyro.plate("a", len(data), subsample_size=subsample_size) as idx:
             assert data[idx].shape == (subsample_size,)