Something went wrong. Try again.
lol
Something went wrong. Try again.
3.6 kB · 137 lines
Nix
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138{ lib, stdenv, buildPythonPackage, fetchFromGitHub,
# build-system setuptools,
# dependencies jax, jaxlib, multipledispatch, numpy, tqdm,
# tests dm-haiku, equinox, flax, funsor, graphviz, optax, pyro-api, pytest-xdist, pytestCheckHook, scikit-learn, tensorflow-probability,}:
buildPythonPackage (finalAttrs: { pname = "numpyro"; version = "0.20.1"; pyproject = true;
src = fetchFromGitHub { owner = "pyro-ppl"; repo = "numpyro"; tag = finalAttrs.version; hash = "sha256-sNqllL9nBwXp0kn+HAjvIaHf7LR0UKh9q7DZ20yCr5A="; };
patches = [ # Remove usage of xla_pmap_p which was removed in jax 0.10.0 ./fix-jax-0.10.0-compat.patch ];
build-system = [ setuptools ];
dependencies = [ jax jaxlib multipledispatch numpy tqdm ];
nativeCheckInputs = [ dm-haiku equinox flax funsor graphviz optax pyro-api pytest-xdist pytestCheckHook scikit-learn tensorflow-probability ];
pythonImportsCheck = [ "numpyro" ];
pytestFlags = [ # Tests memory consumption grows significantly with the number of parallel processes (reaches ~200GB with 80 jobs) "--maxprocesses=8"
# A few tests fail with: # UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. # Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. # You can double-check how many devices are available in your system using `jax.local_device_count()`. "-Wignore::UserWarning"
# FutureWarning: In the future `np.object` will be defined as the corresponding NumPy scalar. "-Wignore::FutureWarning" ];
disabledTests = [ # Failing with jax>=0.9.0 # TypeError: Error interpreting argument to closed_call as a JAX value "test_provenance_call" "test_provenance_closed_call" "test_numpyrooptim_no_double_jit"
# ValueError: Found unexpected Arrays on value of type <class 'list'> in static attribute 'layers' # of Pytree type '<class 'test_module.test_random_nnx_module_mcmc_sequence_params.<locals>.MLP'>'. # This is an error starting from Flax version 0.12.0. "test_random_nnx_module_mcmc_sequence_param"
# AssertionError, assert GLOBAL["count"] == 4 (assert 5 == 4) "test_mcmc_parallel_chain"
# AssertionError due to tolerance issues "test_bijective_transforms" "test_cpu" "test_entropy_categorical" "test_gaussian_model"
# > with pytest.warns(UserWarning, match="Hessian of log posterior"): # E Failed: DID NOT WARN. No warnings of type (<class 'UserWarning'>,) were emitted. # E Emitted warnings: []. "test_laplace_approximation_warning"
# ValueError: compiling computation that requires 2 logical devices, but only 1 XLA devices are available (num_replicas=2) "test_chain"
# Failed: DID NOT RAISE <class 'UserWarning'> "test_interval_censored_validate_sample" ] ++ lib.optionals stdenv.hostPlatform.isDarwin [ # AssertionError: Not equal to tolerance rtol=0.06, atol=0 "test_functional_map" ];
disabledTestPaths = [ # Require internet access "test/test_example_utils.py" ];
meta = { description = "Library for probabilistic programming with NumPy"; homepage = "https://num.pyro.ai/"; changelog = "https://github.com/pyro-ppl/numpyro/releases/tag/${finalAttrs.src.tag}"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ fab ]; };})