Something went wrong. Try again.
What happens if I push all of nixpkgs
Something went wrong. Try again.
2.9 kB · 118 lines
Nix
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119{ lib, stdenv, buildPythonPackage, fetchFromGitHub,
# build-system setuptools,
# dependencies jax, jaxlib, multipledispatch, numpy, tqdm,
# tests dm-haiku, flax, funsor, graphviz, optax, pyro-api, pytest-xdist, pytestCheckHook, scikit-learn, tensorflow-probability,}:
buildPythonPackage rec { pname = "numpyro"; version = "0.18.0"; pyproject = true;
src = fetchFromGitHub { owner = "pyro-ppl"; repo = "numpyro"; tag = version; hash = "sha256-0X/ta2yfzjf3JnZYdUAzQmXvbsDpwFCJe/bArMSWQgU="; };
build-system = [ setuptools ];
dependencies = [ jax jaxlib multipledispatch numpy tqdm ];
nativeCheckInputs = [ dm-haiku 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" ];
disabledTests = [ # 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"
# Failing since flax==0.11.0 # KeyError: "No RngStream named 'dropout' found in Rngs." # https://github.com/pyro-ppl/numpyro/issues/2055 "test_nnx_state_dropout_smoke" ] ++ 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/${version}"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ fab ]; };}