diff --git a/pkgs/development/python-modules/opt-einsum-fx/default.nix b/pkgs/development/python-modules/opt-einsum-fx/default.nix new file mode 100644 index 000000000000..d183970cc557 --- /dev/null +++ b/pkgs/development/python-modules/opt-einsum-fx/default.nix @@ -0,0 +1,52 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + + # dependencies + opt-einsum, + torch, + + # tests + pytestCheckHook, +}: + +buildPythonPackage (finalAttrs: { + pname = "opt-einsum-fx"; + version = "0.1.4"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "Linux-cpp-lisp"; + repo = "opt_einsum_fx"; + tag = "v${finalAttrs.version}"; + hash = "sha256-HamDghqmdX4Q+4zXQvCly588p3TaYFCSnzgEKLVMXSo="; + }; + + build-system = [ + setuptools + ]; + + dependencies = [ + opt-einsum + torch + ]; + + pythonImportsCheck = [ "opt_einsum_fx" ]; + + nativeCheckInputs = [ + pytestCheckHook + ]; + + meta = { + description = "Einsum optimization using opt_einsum and PyTorch FX graph rewriting"; + homepage = "https://github.com/Linux-cpp-lisp/opt_einsum_fx"; + changelog = "https://github.com/Linux-cpp-lisp/opt_einsum_fx/blob/${finalAttrs.src.tag}/CHANGELOG.md"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index cb255c8fb62c..da7abac6b097 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -12052,6 +12052,8 @@ self: super: with self; { opt-einsum = callPackage ../development/python-modules/opt-einsum { }; + opt-einsum-fx = callPackage ../development/python-modules/opt-einsum-fx { }; + optax = callPackage ../development/python-modules/optax { }; optiland = callPackage ../development/python-modules/optiland { };