tjh.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 poetry-core,
8
9 # dependencies
10 jax,
11 jaxlib,
12 tensorflow-probability,
13
14 # tests
15 inference-gym,
16 pytestCheckHook,
17}:
18
19buildPythonPackage rec {
20 pname = "oryx";
21 version = "0.2.7";
22 pyproject = true;
23
24 src = fetchFromGitHub {
25 owner = "jax-ml";
26 repo = "oryx";
27 rev = "refs/tags/v${version}";
28 hash = "sha256-1n7ogGuFNAeOyXWe0/pAouhg2+aA3MXxlCcsrfqRTdU=";
29 };
30
31 build-system = [ poetry-core ];
32
33 dependencies = [
34 jax
35 jaxlib
36 tensorflow-probability
37 ];
38
39 pythonImportsCheck = [ "oryx" ];
40
41 nativeCheckInputs = [
42 inference-gym
43 pytestCheckHook
44 ];
45
46 meta = {
47 description = "Library for probabilistic programming and deep learning built on top of Jax";
48 homepage = "https://github.com/jax-ml/oryx";
49 changelog = "https://github.com/jax-ml/oryx/releases/tag/v${version}";
50 license = lib.licenses.asl20;
51 maintainers = with lib.maintainers; [ GaetanLepage ];
52 };
53}