diff --git a/pkgs/development/python-modules/hyper-connections/default.nix b/pkgs/development/python-modules/hyper-connections/default.nix index 2a10d0d38e9c..b0ec4af38158 100644 --- a/pkgs/development/python-modules/hyper-connections/default.nix +++ b/pkgs/development/python-modules/hyper-connections/default.nix @@ -5,19 +5,21 @@ fetchFromGitHub, hatchling, pytestCheckHook, + stdenv, torch, + torch-einops-utils, }: buildPythonPackage (finalAttrs: { pname = "hyper-connections"; - version = "0.4.7"; + version = "0.4.9"; pyproject = true; src = fetchFromGitHub { owner = "lucidrains"; repo = "hyper-connections"; tag = finalAttrs.version; - hash = "sha256-x1Yx9Fnow9154kFGLmjeCBLYJsbv6oJiC6Rk1XudqJQ="; + hash = "sha256-RDwnRtHUWilyqsDmdiV+kRg7BqTS1yghiu9RAM+MNjE="; }; build-system = [ hatchling ]; @@ -25,10 +27,17 @@ buildPythonPackage (finalAttrs: { dependencies = [ einops torch + torch-einops-utils ]; nativeCheckInputs = [ pytestCheckHook ]; + disabledTests = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [ + # torch's cpuinfo init fails to parse /sys/devices/system/cpu/{possible,present} + # in the build sandbox on aarch64-linux, breaking `.half()` calls + "test_mhc_dtype_restoration" + ]; + pythonImportsCheck = [ "hyper_connections" ]; meta = { diff --git a/pkgs/development/python-modules/torch-einops-utils/default.nix b/pkgs/development/python-modules/torch-einops-utils/default.nix new file mode 100644 index 000000000000..ffb0fd01cb34 --- /dev/null +++ b/pkgs/development/python-modules/torch-einops-utils/default.nix @@ -0,0 +1,41 @@ +{ + lib, + buildPythonPackage, + einops, + fetchFromGitHub, + hatchling, + pytestCheckHook, + torch, +}: + +buildPythonPackage (finalAttrs: { + pname = "torch-einops-utils"; + version = "0.0.29"; + pyproject = true; + + src = fetchFromGitHub { + owner = "lucidrains"; + repo = "torch-einops-utils"; + tag = finalAttrs.version; + hash = "sha256-ja3HeBvAQRyGL2anqIQa2iiHhOZUhF73do7pvrTyRo0="; + }; + + build-system = [ hatchling ]; + + dependencies = [ + einops + torch + ]; + + nativeCheckInputs = [ pytestCheckHook ]; + + pythonImportsCheck = [ "torch_einops_utils" ]; + + meta = { + description = "Utility functions for torch and einops"; + homepage = "https://github.com/lucidrains/torch-einops-utils"; + changelog = "https://github.com/lucidrains/torch-einops-utils/releases/tag/${finalAttrs.src.tag}"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ miniharinn ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index e8558e3a9521..3bcc3092e7dd 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19598,6 +19598,8 @@ self: super: with self; { torch-c-dlpack-ext = callPackage ../development/python-modules/torch-c-dlpack-ext { }; + torch-einops-utils = callPackage ../development/python-modules/torch-einops-utils { }; + torch-geometric = callPackage ../development/python-modules/torch-geometric { }; # Required to test triton