Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677{ lib, buildPythonPackage, fetchFromGitHub, causal-conv1d, einops, ninja, setuptools, torch, transformers, triton, cudaPackages, config, cudaSupport ? config.cudaSupport, which,}:
buildPythonPackage rec { pname = "mamba"; version = "2.2.2"; pyproject = true;
src = fetchFromGitHub { owner = "state-spaces"; repo = "mamba"; tag = "v${version}"; hash = "sha256-R702JjM3AGk7upN7GkNK8u1q4ekMK9fYQkpO6Re45Ng="; };
build-system = [ ninja setuptools torch ];
nativeBuildInputs = [ which ];
buildInputs = ( lib.optionals cudaSupport ( with cudaPackages; [ cuda_cudart # cuda_runtime.h, -lcudart cccl libcusparse # cusparse.h libcusolver # cusolverDn.h cuda_nvcc libcublas ] ) );
dependencies = [ causal-conv1d einops torch transformers triton ];
env = { MAMBA_FORCE_BUILD = "TRUE"; } // lib.optionalAttrs cudaSupport { CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}"; };
# pytest tests not enabled due to nvidia GPU dependency pythonImportsCheck = [ "mamba_ssm" ];
meta = { description = "Linear-Time Sequence Modeling with Selective State Spaces"; homepage = "https://github.com/state-spaces/mamba"; license = lib.licenses.asl20; # The package requires CUDA or ROCm, the ROCm build hasn't # been completed or tested, so broken if not using cuda. broken = !cudaSupport; };}