diff --git a/pkgs/development/python-modules/pytorchviz/default.nix b/pkgs/development/python-modules/pytorchviz/default.nix new file mode 100644 index 000000000000..445e2b9b425d --- /dev/null +++ b/pkgs/development/python-modules/pytorchviz/default.nix @@ -0,0 +1,46 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + distutils, + fsspec, + graphviz, + torch, + unittestCheckHook, +}: + +buildPythonPackage rec { + pname = "pytorchviz"; + version = "0.0.2-unstable-2024-12-30"; + pyproject = true; + + src = fetchFromGitHub { + owner = "szagoruyko"; + repo = "pytorchviz"; + # No tags in the upstream GitHub repo + rev = "5cf04c13e601366f6b9cf5939b5af5144d55b887"; + hash = "sha256-La1X8Y64n/vNGDUEsw1iZ5Mb6/w3WayeWxa62QxLyHA="; + }; + + dependencies = [ + graphviz + torch + ]; + + nativeCheckInputs = [ + unittestCheckHook + distutils + fsspec + ]; + + unittestFlagsArray = [ "test/" ]; + + pythonImportsCheck = [ "torchviz" ]; + + meta = { + description = "Small package to create visualizations of PyTorch execution graphs"; + homepage = "https://github.com/szagoruyko/pytorchviz"; + license = lib.licenses.mit; + maintainers = with lib.maintainers; [ jherland ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 012648099650..0edf8fd028c6 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -15237,6 +15237,8 @@ self: super: with self; { pytorch3d = callPackage ../development/python-modules/pytorch3d { }; + pytorchviz = callPackage ../development/python-modules/pytorchviz { }; + pytouchline-extended = callPackage ../development/python-modules/pytouchline-extended { }; pytouchlinesl = callPackage ../development/python-modules/pytouchlinesl { };