python3Packages.objax: init at 1.4.0

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

@@ -0,0 +1,46 @@
+{ lib
+, fetchFromGitHub
+, buildPythonPackage
+, jax
+, jaxlib
+, numpy
+, parameterized
+, pillow
+, scipy
+, tensorflow-tensorboard_2 ? null
+}:
+
+buildPythonPackage rec {
+  pname = "objax";
+  version = "1.4.0";
+
+  src = fetchFromGitHub {
+    owner = "google";
+    repo = "objax";
+    rev = "v${version}";
+    sha256 = "09gm61ghn5mi92q5mhx22mcv6aa6z78jsrnfar1hd3nwwyn9dq42";
+  };
+
+  propagatedBuildInputs = [
+    jax
+    jaxlib
+    numpy
+    parameterized
+    pillow
+    scipy
+    tensorflow-tensorboard_2
+  ];
+
+  pythonImportsCheck = [
+    "objax"
+  ];
+
+  meta = with lib; {
+    description = "Objax is a machine learning framework that provides an Object Oriented layer for JAX.";
+    homepage = "https://github.com/google/objax";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+    # Darwin doesn't have `tensorflow-tensorboard_2` which is required by wheel deps.
+    platforms = [ "aarch64-linux" "x86_64-linux" ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -5460,6 +5460,8 @@ in {
 
   obfsproxy = callPackage ../development/python-modules/obfsproxy { };
 
+  objax = callPackage ../development/python-modules/objax { };
+
   objgraph = callPackage ../development/python-modules/objgraph {
     # requires both the graphviz package and python package
     graphvizPkgs = pkgs.graphviz;