mypy

[vineetbansal]

If we are to leverage the use of adding type-hints (as you've done in the functions that have been doc-stringed), we should add this to the pre-commit config:

  - repo: https://github.com/pre-commit/mirrors-mypy
    rev: "v1.13.0"
    hooks:
      - id: mypy
        files: "(src|tests)"

This will likely uncover places where the types are incorrect. I notice, for example:

def stack_slices_center(
   center_slice: AnnData,
   slices: list[AnnData],
   pis: list[np.ndarray],
   ...

While I'm pretty sure pis is a torch.Tensor (if not, the calling function is probably mis-typed as well).