Commit 534cc954 authored by Madoura's avatar Madoura Committed by Weijia Wang
Browse files

python3Packages.torch: Add 'tritonSupport' option 

Guard this option against 'stdenv.isDarwin', which currently is unknown whether it works or not.
parent da817841

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

+3 −4
Original line number Diff line number Diff line
@@ -16,9 +16,10 @@ 
  filelock,

  jinja2,

  networkx,

  openai-triton,

  sympy,

  numpy, pyyaml, cffi, click, typing-extensions,

  # ROCm build and `torch.compile` requires `openai-triton`

  tritonSupport ? (!stdenv.isDarwin), openai-triton,



  # Unit tests

  hypothesis, psutil,
@@ -377,12 +378,10 @@ in buildPythonPackage rec { 
    # the following are required for tensorboard support

    pillow six future tensorboard protobuf



    # ROCm build and `torch.compile` requires openai-triton

    openai-triton



    # torch/csrc requires `pybind11` at runtime

    pybind11

  ]

  ++ lib.optionals tritonSupport [ openai-triton ]

  ++ lib.optionals MPISupport [ mpi ]

  ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];