Something went wrong. Try again.
What happens if I push all of nixpkgs
Something went wrong. Try again.
1.6 kB · 78 lines
Nix
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879{ lib, buildPythonPackage, fetchFromGitHub, causal-conv1d, einops, ninja, setuptools, torch, transformers, triton, cudaPackages, rocmPackages, 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 cuda_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 = with lib; { description = "Linear-Time Sequence Modeling with Selective State Spaces"; homepage = "https://github.com/state-spaces/mamba"; license = licenses.asl20; maintainers = with maintainers; [ cfhammill ]; # The package requires CUDA or ROCm, the ROCm build hasn't # been completed or tested, so broken if not using cuda. broken = !cudaSupport; };}