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