Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283{ lib, buildPythonPackage, fetchFromGitHub,
# build-system hatchling,
# dependencies jax, jaxtyping, typing-extensions, wadler-lindig,
# tests beartype, optax, pytest-xdist, pytestCheckHook,}:
buildPythonPackage (finalAttrs: { pname = "equinox"; version = "0.13.8"; pyproject = true; __structuredAttrs = true;
src = fetchFromGitHub { owner = "patrick-kidger"; repo = "equinox"; tag = "v${finalAttrs.version}"; hash = "sha256-JiIZKWuSkvrF09GdmegUeTyidaM5IRp4uqjJRsn86E4="; };
# Relax speed constraints on tests that can fail on busy builders postPatch = '' substituteInPlace tests/test_while_loop.py \ --replace-fail "speed < 0.1" "speed < 0.5" \ --replace-fail "speed < 0.5" "speed < 1" \ --replace-fail "speed < 1" "speed < 20" \ --replace-fail "speed < 2" "speed < 20" '';
build-system = [ hatchling ];
dependencies = [ jax jaxtyping typing-extensions wadler-lindig ];
nativeCheckInputs = [ beartype optax pytest-xdist pytestCheckHook ];
pythonImportsCheck = [ "equinox" ];
pytestFlags = [ # UserWarning: Explicitly requested dtype int64 requested in ShapeDtypeStruct is not available, and will be truncated to dtype int32. # To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. # See https://github.com/jax-ml/jax#current-gotchas for more. "-Wignore::UserWarning" ];
disabledTests = [ # Flaky under heavy load: # AssertionError: Non-linear scaling detected: ratio=1.56 "test_speed_buffer_while" ];
meta = { description = "JAX library based around a simple idea: represent parameterised functions (such as neural networks) as PyTrees"; changelog = "https://github.com/patrick-kidger/equinox/releases/tag/${finalAttrs.src.tag}"; homepage = "https://github.com/patrick-kidger/equinox"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ GaetanLepage ]; };})