Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103{ lib, buildPythonPackage,
# build-system setuptools,
# dependencies fla-core, transformers,
# optional-dependencies tilelang, causal-conv1d, pandas, matplotlib, datasets,
# tests pytestCheckHook, writableTmpDirAsHomeHook,
# passthru flash-linear-attention,}:
buildPythonPackage (finalAttrs: { pname = "flash-linear-attention"; inherit (fla-core) version src; pyproject = true; __structuredAttrs = true;
# Only ships fla/{layers,models}. The rest of `fla` is provided by fla-core. postPatch = '' python -c "import sys; sys.path.insert(0, 'scripts'); import build_packages; build_packages.build_split_packages('split')" cd split/flash-linear-attention '';
build-system = [ setuptools ];
dependencies = [ fla-core transformers ];
optional-dependencies = { tilelang = [ tilelang ]; conv1d = [ causal-conv1d ]; benchmark = [ matplotlib pandas datasets ]; };
pythonImportsCheck = [ "fla" "fla.layers" "fla.models" ];
nativeCheckInputs = [ pytestCheckHook # torch._inductor.exc.InductorError: PermissionError: [Errno 13] Permission denied: '/homeless-shelter' writableTmpDirAsHomeHook ];
# Run the tests from the repo root against the installed package preCheck = '' cd "$NIX_BUILD_TOP/$sourceRoot" rm -rf fla split '';
# The full test suite takes hours: only run a few representative tests. enabledTestPaths = [ "tests/layers/test_gated_deltanet.py" "tests/models/test_cache.py::test_crop" "tests/models/test_cache.py::test_pickle" "tests/models/test_cache.py::test_reset" "tests/models/test_generation.py::test_generate" "tests/models/test_modeling_gated_deltanet.py::test_generation" "tests/models/test_modeling_gated_deltanet.py::test_modeling" "tests/models/test_modeling_gla.py::test_generation" "tests/modules/test_conv.py::test_conv" ];
# Tests require a GPU doCheck = false; passthru.gpuCheck = flash-linear-attention.overridePythonAttrs { requiredSystemFeatures = [ "cuda" ]; doCheck = true; };
meta = { description = "Triton-based implementations of causal linear attention"; homepage = "https://github.com/fla-org/flash-linear-attention"; changelog = "https://github.com/fla-org/flash-linear-attention/releases/tag/${finalAttrs.src.tag}"; license = lib.licenses.mit; maintainers = with lib.maintainers; [ BatteredBunny ]; teams = [ lib.teams.cuda ]; };})