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

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchFromGitHub,
   5
   6   # build-system
   7   hatchling,
   8
   9   # dependencies
  10   jax,
  11   jaxlib,
  12   jaxtyping,
  13   typing-extensions,
  14
  15   # checks
  16   beartype,
  17   optax,
  18   pytest-xdist,
  19   pytestCheckHook,
  20 }:
  21
  22 buildPythonPackage rec {
  23   pname = "equinox";
  24   version = "0.11.8";
  25   pyproject = true;
  26
  27   src = fetchFromGitHub {
  28     owner = "patrick-kidger";
  29     repo = "equinox";
  30     rev = "refs/tags/v${version}";
  31     hash = "sha256-lZb2NobSELz8kviPd4Z8PPEEaydaEC5Z6eb9pzC7Ki8=";
  32   };
  33
  34   build-system = [ hatchling ];
  35
  36   dependencies = [
  37     jax
  38     jaxlib
  39     jaxtyping
  40     typing-extensions
  41   ];
  42
  43   nativeCheckInputs = [
  44     beartype
  45     optax
  46     pytest-xdist
  47     pytestCheckHook
  48   ];
  49
  50   pythonImportsCheck = [ "equinox" ];
  51
  52   meta = {
  53     description = "JAX library based around a simple idea: represent parameterised functions (such as neural networks) as PyTrees";
  54     changelog = "https://github.com/patrick-kidger/equinox/releases/tag/v${version}";
  55     homepage = "https://github.com/patrick-kidger/equinox";
  56     license = lib.licenses.asl20;
  57     maintainers = with lib.maintainers; [ GaetanLepage ];
  58   };
  59 }