python312Packages.torch-geometric: init at 2.6.1

{

  lib,

  buildPythonPackage,

  fetchFromGitHub,

  # build-system

  flit-core,

  # dependencies

  aiohttp,

  fsspec,

  jinja2,

  numpy,

  psutil,

  pyparsing,

  requests,

  torch,

  tqdm,

  # optional-dependencies

  matplotlib,

  networkx,

  pandas,

  protobuf,

  wandb,

  ipython,

  matplotlib-inline,

  pre-commit,

  torch-geometric,

  ase,

  # captum,

  graphviz,

  h5py,

  numba,

  opt-einsum,

  pgmpy,

  pynndescent,

  # pytorch-memlab,

  rdflib,

  rdkit,

  scikit-image,

  scikit-learn,

  scipy,

  statsmodels,

  sympy,

  tabulate,

  torchmetrics,

  trimesh,

  pytorch-lightning,

  yacs,

  huggingface-hub,

  onnx,

  onnxruntime,

  pytest,

  pytest-cov,

  # tests

  pytestCheckHook,

}:

buildPythonPackage rec {

  pname = "torch-geometric";

  version = "2.6.1";

  pyproject = true;

  src = fetchFromGitHub {

    owner = "pyg-team";

    repo = "pytorch_geometric";

    rev = "refs/tags/${version}";

    hash = "sha256-Zw9YqPQw2N0ZKn5i5Kl4Cjk9JDTmvZmyO/VvIVr6fTU=";

  };

  build-system = [

    flit-core

  ];

  dependencies = [

    aiohttp

    fsspec

    jinja2

    numpy

    psutil

    pyparsing

    requests

    torch

    tqdm

  ];

  optional-dependencies = {

    benchmark = [

      matplotlib

      networkx

      pandas

      protobuf

      wandb

    ];

    dev = [

      ipython

      matplotlib-inline

      pre-commit

      torch-geometric

    ];

    full = [

      ase

      # captum

      graphviz

      h5py

      matplotlib

      networkx

      numba

      opt-einsum

      pandas

      pgmpy

      pynndescent

      # pytorch-memlab

      rdflib

      rdkit

      scikit-image

      scikit-learn

      scipy

      statsmodels

      sympy

      tabulate

      torch-geometric

      torchmetrics

      trimesh

    ];

    graphgym = [

      protobuf

      pytorch-lightning

      yacs

    ];

    modelhub = [

      huggingface-hub

    ];

    test = [

      onnx

      onnxruntime

      pytest

      pytest-cov

    ];

  };

  pythonImportsCheck = [

    "torch_geometric"

  ];

  nativeCheckInputs = [

    pytestCheckHook

  ];

  preCheck = ''

    export HOME=$(mktemp -d)

  '';

  disabledTests = [

    # TODO: try to re-enable when triton will have been updated to 3.0

    # torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:

    # LoweringException: ImportError: cannot import name 'triton_key' from 'triton.compiler.compiler'

    "test_compile_hetero_conv_graph_breaks"

    "test_compile_multi_aggr_sage_conv"

    # RuntimeError: addmm: computation on CPU is not implemented for SparseCsr + SparseCsr @ SparseCsr without MKL.

    # PyTorch built with MKL has better support for addmm with sparse CPU tensors.

    "test_asap"

    "test_graph_unet"

    # AttributeError: type object 'Any' has no attribute '_name'

    "test_type_repr"

  ];

  meta = {

    description = "Graph Neural Network Library for PyTorch";

    homepage = "https://github.com/pyg-team/pytorch_geometric";

    changelog = "https://github.com/pyg-team/pytorch_geometric/blob/${src.rev}/CHANGELOG.md";

    license = lib.licenses.mit;

    maintainers = with lib.maintainers; [ GaetanLepage ];

  };

}

  torch-audiomentations = callPackage ../development/python-modules/torch-audiomentations { };

  torch-geometric = callPackage ../development/python-modules/torch-geometric { };

  torch-pitch-shift = callPackage ../development/python-modules/torch-pitch-shift { };

  torch-bin = callPackage ../development/python-modules/torch/bin.nix {