From 8f8a5959cced15a23aef1f388d3b74d205dbfa80 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 30 Mar 2023 18:34:27 +0200
Subject: [PATCH] python3Packages.jax-jumpy: init at 1.0.0
---
.../python-modules/jax-jumpy/default.nix | 33 +++++++++++++++++++
pkgs/top-level/python-packages.nix | 2 ++
2 files changed, 35 insertions(+)
create mode 100644 pkgs/development/python-modules/jax-jumpy/default.nix
diff --git a/pkgs/development/python-modules/jax-jumpy/default.nix b/pkgs/development/python-modules/jax-jumpy/default.nix
new file mode 100644
index 000000000000..f37677da0283
--- /dev/null
+++ b/pkgs/development/python-modules/jax-jumpy/default.nix
@@ -0,0 +1,33 @@
+{ lib
+, buildPythonPackage
+, fetchFromGitHub
+, numpy
+, setuptools
+}:
+
+buildPythonPackage rec {
+ pname = "jumpy";
+ version = "1.0.0";
+
+ src = fetchFromGitHub {
+ owner = "Farama-Foundation";
+ repo = pname;
+ rev = version;
+ hash = "sha256-tPQ/v2AVnAEC+08BVAvvgJ8Pj89nXZSn2tQ6nxXuSfA=";
+ };
+
+ format = "pyproject";
+
+ nativeBuildInputs = [ setuptools ];
+
+ propagatedBuildInputs = [ numpy ];
+
+ pythonImportsCheck = [ "jumpy" ];
+
+ meta = with lib; {
+ description = "Jumpy is a common backend for NumPy and optionally JAX";
+ homepage = "https://github.com/Farama-Foundation/Jumpy";
+ license = licenses.asl20;
+ maintainers = with maintainers; [ GaetanLepage ];
+ };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index a1a8dd080fca..53e7aa1048b3 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -4922,6 +4922,8 @@ self: super: with self; {
jax = callPackage ../development/python-modules/jax { };
+ jax-jumpy = callPackage ../development/python-modules/jax-jumpy { };
+
jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix {
cudaSupport = pkgs.config.cudaSupport or false;
};