pkgs/development/python-modules/oryx/default.nix

   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
  19 buildPythonPackage 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 }