tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 setuptools,
8
9 # dependencies
10 dotmap,
11 flax,
12 jax,
13 matplotlib,
14 numpy,
15
16 # tests
17 # brax, (unpackaged)
18 # gymnax, (unpackaged)
19 pytestCheckHook,
20 torch,
21 torchvision,
22 writableTmpDirAsHomeHook,
23}:
24
25buildPythonPackage rec {
26 pname = "evosax";
27 version = "0.2.0";
28 pyproject = true;
29
30 src = fetchFromGitHub {
31 owner = "RobertTLange";
32 repo = "evosax";
33 tag = "v.${version}";
34 hash = "sha256-ye5IHM8Pn/+BXI9kcB3W281Gna9hXV8DwsaJ9Xu06fU=";
35 };
36
37 build-system = [ setuptools ];
38
39 dependencies = [
40 dotmap
41 flax
42 jax
43 matplotlib
44 numpy
45 ];
46
47 pythonImportsCheck = [ "evosax" ];
48
49 nativeCheckInputs = [
50 # brax
51 # gymnax
52 pytestCheckHook
53 torch
54 torchvision
55 writableTmpDirAsHomeHook
56 ];
57
58 disabledTests = [
59 # Requires unpackaged gymnax
60 "test_env_ffw_rollout"
61
62 # TypeError: ShapedArray.__init__() got an unexpected keyword argument 'named_shape'
63 "test_base_api"
64 "test_run"
65 "test_run_scan"
66
67 # Tries to download a data set from the internet
68 "test_brax_problem_eval"
69 "test_brax_problem_init"
70 "test_brax_problem_sample"
71 "test_gymnax_problem_eval"
72 "test_gymnax_problem_init"
73 "test_gymnax_problem_sample"
74 "test_torchvision_problem_eval"
75 "test_torchvision_problem_init"
76 "test_torchvision_problem_sample"
77 "test_vision_fitness"
78 ];
79
80 meta = {
81 description = "Evolution Strategies in JAX";
82 homepage = "https://github.com/RobertTLange/evosax";
83 changelog = "https://github.com/RobertTLange/evosax/releases/tag/v.${version}";
84 license = lib.licenses.asl20;
85 maintainers = with lib.maintainers; [ GaetanLepage ];
86 };
87}