Add batch_converter hook + multiprocessing support to experimental.pytorch.AnnLoader

[ronamit]

Add batch_converter hook + multiprocessing support to experimental.pytorch.AnnLoader

This PR solves two major AnnLoader limitations:

1. Batch-level transformation: Add batch_converter parameter for advanced post-processing
2. Multiprocessing support: Enable num_workers > 0 (previously crashed due to unpicklable AnnCollectionView)

Motivation

AnnLoader currently offers only element-wise converters via convert["X"] mapping, insufficient for:

• Returning dict instead of AnnCollectionView
• On-the-fly augmentations using both .X and .obs
• PyTorch Lightning integration with specific batch signatures
• Adding derived fields or metadata to batches

Additionally, num_workers > 0 has never worked, preventing parallel data loading in production workflows.

Implementation

1. batch_converter parameter

AnnLoader(adata, batch_size=128, batch_converter=my_fn, num_workers=4)

Optional callable applied to each batch before returning to user. Fully backward-compatible.

2. Multiprocessing support

• __getstate__/__setstate__ hooks enable AnnCollectionView pickling across worker processes
• Worker-side batch conversion via custom collate_fn
• Per-worker HDF5 file handles following h5py best practices

3. Helper converter

batch_dict_converter converts AnnCollectionView → dict[str, Tensor] with "x" key for .X and keys for each .obs column.

Usage

from anndata.experimental.pytorch import AnnLoader, batch_dict_converter

# Single-threaded
loader = AnnLoader(adata, batch_size=256, batch_converter=batch_dict_converter, num_workers=0)

# Multi-threaded (now possible)
loader = AnnLoader(adata, batch_size=256, batch_converter=batch_dict_converter, num_workers=4)

for batch in loader:
    x = batch["x"]          # torch.Tensor
    obs_fields = batch["cell_type"]  # obs columns as tensors

Performance

scenario	master	this PR
Single-threaded	133 ms	134 ms (+0.8%)
Multiprocessing	Crashes	Works

Negligible overhead for single-threaded use. Primary benefit is enabling parallel data loading for I/O-bound workflows.

Testing

• Comprehensive unit tests for both single and multi-threaded modes
• Validates backward compatibility (no converter → AnnCollectionView)
• Documents multiprocessing necessity vs vanilla DataLoader

---

Backward compatible: Optional parameter defaults to None, no breaking changes.