Commit Graph
1 Commits
Author SHA1 Message Date
Omkar Kabde df89c133ca feat(python)!: align Permutation.with_format("torch") with HuggingFace set_format("torch") (#3369)
Closes #3245.

> **BREAKING CHANGE:** `with_format("torch")` no longer returns a list
of stacked row tensors. It now returns per-row dicts so PyTorch's
default `DataLoader` collate stacks them into `{col: tensor(B,)}`.
Switch to `with_format("torch_row")` to keep the old shape.

### What changed

`"torch"` now returns a list of per-row dicts (`[{col: tensor}, ...]`)
at every indexed access path. The default `DataLoader` collate stacks
them into a column-keyed batched dict, no custom `collate_fn` needed.
The old shape is preserved under a new `"torch_row"` literal.
`"torch_col"` is unchanged.

The unbatching lives inside the transform (`batch_to_tensor_dict`), not
`__getitems__`, so the shape survives pickling and works under
`DataLoader(num_workers>0, multiprocessing_context="spawn")`.

### Format comparison

| Format | `iter(batch_size=N)` | `__getitems__([0,1,2])` | `DataLoader`
default collate |
|---|---|---|---|
| `"torch"` (new) | `list[{col: tensor}]` length N | `list[{col:
tensor}]` length 3 | `{col: tensor(B,)}` |
| `"torch_row"` (old `"torch"` behavior) | `list[tensor(n_cols,)]`
length N | `list[tensor(n_cols,)]` length 3 | `tensor(B, n_cols)` |
| `"torch_col"` (unchanged) | `tensor(n_cols, N)` | `tensor(n_cols, 3)`
| needs `collate_fn=lambda x: x` |

Output matches HuggingFace `Dataset.set_format("torch")` on container
shape, keys, and values at every access path. The only divergence:
HuggingFace downcasts `float64` to `torch.float32` by default, LanceDB
preserves dtype. Verified by `scripts/verify_torch_format.py`.

### Migration

```python
# Old default — column names lost, shape was tensor(B, n_cols)
DataLoader(Permutation.identity(table).with_format("torch"))

# New default — column names preserved
DataLoader(Permutation.identity(table).with_format("torch"))     # {col: tensor(B,)}

# Keep old behavior
DataLoader(Permutation.identity(table).with_format("torch_row")) # tensor(B, n_cols)
```
2026-07-07 15:13:09 -07:00