[ma-sadeghi]

Why

Since the backend wrapper (27bbfeb), xlb>=0.3.1 and warp-lang==1.10.0 are in [project].dependencies, so every install of kabs pulls in XLB, JAX/jaxlib, Warp and their dependencies (ml-dtypes, opt-einsum, numpy-stl, python-utils, plus pydantic and trimesh through XLB). Only the XLB backend uses them: kabs/_solve_flow_xlb.py and kabs/_warp_convergence.py, and their imports were already lazy.

Downstream projects that only use the Taichi solver pay for that on every install. In PoreStore, which installs kabs in its development and CI environments, the XLB stack is about 340 MB of the environment on macOS arm64 (jaxlib alone is about 256 MB, Warp 65 MB) and none of it is ever imported. In kabs's own environment on Linux x86_64 the gap is larger: 1,942 MB with the XLB stack against 1,085 MB without it.

What changes

• pyproject.toml: xlb and warp-lang move to a new optional extra, xlb. Install it with pip install "kabs[xlb]", or uv sync --extra xlb in a clone. The pins are unchanged. uv.lock is regenerated with uv 0.8.17, the version CI uses; the only change is that the two packages now sit under the extra.
• kabs/_solve_flow_xlb.py: selecting the XLB backend without the extra (solve_flow(backend="xlb") or solve_flow_xlb(...)) raises an ImportError that names kabs[xlb] and the uv sync --extra xlb command, chained to the original ModuleNotFoundError. Input validation still runs before the import, so a bad direction or compute_backend is still reported as a ValueError without XLB. The docstring's installation note points at the extra.
• Imports: import kabs loads none of xlb, jax or warp (this was already true, and a new test now enforces it). jax is imported only inside the JAX convergence helpers, xlb only inside solve_flow_xlb after validation, and warp only on the compute_backend="warp" path through kabs._warp_convergence, which nothing imports at package level.
• README.md and docs/warp-backend-readiness.md: the install instructions name the extra.

Tests

• tests/test_convergence_xlb.py skips without xlb (pytest.importorskip), as tests/test_solve_flow_xlb.py already did. It used to import jax at module level, which would fail collection without the extra.
• tests/test_backend_import_isolation.py::test_import_kabs_does_not_import_the_xlb_stack checks, in a fresh process, that import kabs does not load xlb, jax, warp or kabs._warp_convergence.
• tests/test_solve_flow_xlb_validation.py::test_missing_xlb_extra_is_named_in_the_error[jax|warp] hides xlb and checks the error through the public solve_flow. It runs in both environments.
• tests/test_solve_flow_xlb_validation.py::test_install_without_the_extra_omits_the_xlb_stack runs when KABS_TEST_WITHOUT_XLB=1, the same opt-in pattern as KABS_TEST_WARP_CUDA. It checks that xlb, warp and jax are really absent and that solve_flow(backend="xlb") raises the named error.
• .github/workflows/tests.yml runs the suite twice: once with uv sync --group dev --extra xlb, which runs every test and uploads coverage as before, and once with uv sync --group dev and KABS_TEST_WITHOUT_XLB=1. Tests run with uv run --no-sync, so the job's environment is exactly what its sync step installed.

Results on Linux x86_64 (Python 3.12, CPU only, so the KABS_TEST_WARP_CUDA tests skip in every run):

Environment	Result
dev at 2f6d82f (before)	216 passed, 17 skipped
This branch, --extra xlb	219 passed, 18 skipped (the 3 new tests pass; the skipped one is the without-extra check)
This branch, no extra, KABS_TEST_WITHOUT_XLB=1	178 passed, 17 skipped (the XLB modules skip; the missing-extra tests pass)

uvx ruff@0.15.6 check passes on the changed Python files.

One thing I noticed while testing, unrelated to this change: on macOS arm64, tests/test_convergence_xlb.py::test_warp_cpu_combined_reduction_matches_masked_numpy aborts the whole pytest process (Fatal Python error: Aborted) inside Warp's library load (warp/_src/context.py, load_dll). The session fixture in tests/conftest.py has already initialised Taichi by then; I did not check whether that is the cause. It passes on Linux, so CI is unaffected, and a macOS install without the extra skips it.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.0%. Comparing base (2f6d82f) to head (30573e5).

Additional details and impacted files

@@           Coverage Diff           @@
##             dev     #15     +/-   ##
=======================================
+ Coverage   87.8%   88.0%   +0.1%     
=======================================
  Files         11      11             
  Lines       1300    1300             
=======================================
+ Hits        1142    1144      +2     
+ Misses       158     156      -2     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.