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,)