tensorplay.func.chunk_vmap

tensorplay.func.chunk_vmap(func: Callable, in_dims: int | tuple = 0, out_dims: int | tuple[int, ...] | None = 0, randomness: str = 'error', chunks: int = 2) → Callable

vmap() splitting the batch into chunks pieces.

Prefer vmap(..., chunk_size=...), which states the bound in samples rather than in pieces and so does not change meaning with the batch size.