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