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

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchFromGitHub,
   5
   6   # build-system
   7   setuptools,
   8
   9   # dependencies
  10   absl-py,
  11   jax,
  12   jaxlib,
  13   numpy,
  14   toolz,
  15   typing-extensions,
  16
  17   # tests
  18   cloudpickle,
  19   dm-tree,
  20   pytestCheckHook,
  21   pythonOlder,
  22 }:
  23
  24 buildPythonPackage rec {
  25   pname = "chex";
  26   version = "0.1.88";
  27   pyproject = true;
  28
  29   src = fetchFromGitHub {
  30     owner = "deepmind";
  31     repo = "chex";
  32     tag = "v${version}";
  33     hash = "sha256-umRq+FZwyx1hz839ZibRTEFKjbBugrfUJuE8PagjqI4=";
  34   };
  35
  36   build-system = [ setuptools ];
  37
  38   dependencies = [
  39     absl-py
  40     jax
  41     jaxlib
  42     numpy
  43     toolz
  44     typing-extensions
  45   ];
  46
  47   pythonImportsCheck = [ "chex" ];
  48
  49   nativeCheckInputs = [
  50     cloudpickle
  51     dm-tree
  52     pytestCheckHook
  53   ];
  54
  55   # AttributeError: module 'unittest' has no attribute 'makeSuite'
  56   # https://github.com/google-deepmind/chex/issues/371
  57   # TODO: re-enable at next release
  58   doCheck = pythonOlder "3.13";
  59
  60   meta = {
  61     description = "Library of utilities for helping to write reliable JAX code";
  62     homepage = "https://github.com/deepmind/chex";
  63     changelog = "https://github.com/google-deepmind/chex/releases/tag/v${version}";
  64     license = lib.licenses.asl20;
  65     maintainers = with lib.maintainers; [ ndl ];
  66   };
  67 }