Provide tensor-friendly methods for torch and jax through OMEArrow.to_dlpack

.github/workflows/run-tests.yml

@@ -28,7 +28,7 @@ jobs:
   run_tests:
     strategy:
       matrix:
-        python_version: ["3.10", "3.11", "3.12", "3.13"]
+        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
         os: [ubuntu-24.04, macos-14]
     runs-on: ${{ matrix.os }}
     env:
@@ -42,7 +42,7 @@ jobs:
           python-version: ${{ matrix.python_version }}
       - name: Install the latest version of uv
         uses: astral-sh/setup-uv@v7
-      - name: Sync dependencies (with viz extra)
-        run: uv sync --frozen --extra viz
+      - name: Sync dependencies (with viz + dlpack extras)
+        run: uv sync --frozen --extra viz --extra dlpack
       - name: Run pytest
         run: uv run --frozen pytest

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
     -   id: check-yaml
     -   id: detect-private-key
 -   repo: https://github.com/tox-dev/pyproject-fmt
-    rev: "v2.15.3"
+    rev: "v2.16.1"
     hooks:
     -   id: pyproject-fmt
 -   repo: https://github.com/citation-file-format/cffconvert
@@ -39,12 +39,12 @@ repos:
     -   id: yamllint
         exclude: pre-commit-config.yaml
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.15.0"
+    rev: "v0.15.1"
     hooks:
     -   id: ruff-format
     -   id: ruff-check
 -   repo: https://github.com/rhysd/actionlint
-    rev: v1.7.10
+    rev: v1.7.11
     hooks:
     -   id: actionlint
 -   repo: https://gitlab.com/vojko.pribudic.foss/pre-commit-update

README.md

@@ -77,6 +77,33 @@ oa_image.export(how="ome-parquet", out="your_image.ome.parquet")
 oa_image.export(how="vortex", out="your_image.vortex")
 ```
 
+## Tensor view (DLPack)
+
+For tensor-focused workflows (PyTorch/JAX), use `tensor_view` and DLPack export.
+
+```python
+from ome_arrow import OMEArrow
+
+oa = OMEArrow("your_image.ome.parquet")
+
+# Spatial ROI per plane
+view = oa.tensor_view(t=0, z=0, roi=(32, 32, 128, 128), layout="CHW")
+
+# Convenience 3D ROI (x, y, z, w, h, d)
+view3d = oa.tensor_view(roi3d=(32, 32, 2, 128, 128, 4), layout="TZCHW")
+
+# 3D tiled iteration over (z, y, x)
+for cap in view3d.iter_tiles_3d(tile_size=(2, 64, 64), mode="numpy"):
+    pass
+```
+
+Advanced options:
+
+- `chunk_policy="auto" | "combine" | "keep"` controls ChunkedArray handling.
+- `channel_policy="error" | "first"` controls behavior when dropping `C` from layout.
+
+See full docs: [`docs/src/dlpack.md`](docs/src/dlpack.md)
+
 ## Contributing, Development, and Testing
 
 Please see our [contributing documentation](https://github.com/wayscience/ome-arrow/tree/main/CONTRIBUTING.md) for more details on contributions, development, and testing.

docs/src/dlpack.md

@@ -0,0 +1,118 @@
+# Exporting OME-Arrow pixel data via DLPack
+
+OME-Arrow exposes a small tensor view API for pixel data. The returned
+`TensorView` can export DLPack capsules for zero-copy interoperability on CPU
+and (optionally) GPU.
+
+Key defaults:
+
+- OME-Arrow tensor layouts always include channels (`C`) as a tensor axis.
+- Default layout is `CHW` when both `T` and `Z` are singleton in the source.
+- Otherwise, default layout is `TZCHW` (with singleton `T`/`Z` retained unless you override layout).
+- You can override with any valid TZCHW permutation/subset, for example `HWC`, `ZCHW`, or `CHW`.
+
+Layout nomenclature:
+
+- `T`: time index
+- `Z`: z/depth index
+- `C`: channel index
+- `H`: image height (Y axis)
+- `W`: image width (X axis)
+
+Practical mapping:
+
+- 2D image content (`YX`) is typically exposed as `CHW`.
+- 3D z-stack content (`ZYX`) is typically exposed as `ZCHW` or `TZCHW` (with `T=1`).
+- Time-lapse and volumetric content use `TZCHW` by default.
+
+## PyTorch
+
+```python
+from ome_arrow import OMEArrow
+
+obj = OMEArrow("example.ome.parquet")
+view = obj.tensor_view(t=0, z=0, c=0)
+
+# DLPack capsule -> torch.Tensor
+import torch
+
+capsule = view.to_dlpack(mode="arrow", device="cpu")
+flat = torch.utils.dlpack.from_dlpack(capsule)
+tensor = flat.reshape(view.shape)
+```
+
+## JAX
+
+```python
+from ome_arrow import OMEArrow
+
+obj = OMEArrow("example.ome.parquet")
+view = obj.tensor_view(t=0, z=0, c=0, layout="CHW")
+
+import jax.numpy as jnp
+
+capsule = view.to_dlpack(mode="arrow", device="cpu")
+flat = jnp.from_dlpack(capsule)
+arr = flat.reshape(view.shape)
+```
+
+## Iteration examples
+
+```python
+from ome_arrow import OMEArrow
+import numpy as np
+
+obj = OMEArrow("example.ome.parquet")
+view = obj.tensor_view()
+
+# Batch over time (T) dimension.
+for cap in view.iter_dlpack(batch_size=2, shuffle=False, mode="numpy"):
+    batch = np.from_dlpack(cap)
+    # batch shape: (batch, Z, C, H, W) in TZCHW layout
+```
+
+```python
+from ome_arrow import OMEArrow
+import numpy as np
+
+obj = OMEArrow("example.ome.parquet")
+view = obj.tensor_view(t=0, z=0)
+
+# Tile over spatial region.
+for cap in view.iter_dlpack(
+    tile_size=(256, 256), shuffle=True, seed=123, mode="numpy"
+):
+    tile = np.from_dlpack(cap)
+    # tile shape: (C, H, W) in CHW layout
+```
+
+## Ownership and lifetime
+
+`TensorView.to_dlpack()` returns a DLPack-capable object (with `__dlpack__`)
+that references the underlying Arrow values buffer in `mode="arrow"`, or a
+NumPy buffer in `mode="numpy"`. Keep the `TensorView` (or any NumPy array
+returned by `to_numpy`) alive until the consumer finishes using the DLPack
+object.
+
+`mode="arrow"` currently requires a single `(t, z, c)` selection and a full-frame
+ROI. Use `mode="numpy"` for batches, crops, or layout reshaping beyond a simple
+reshape.
+
+Zero-copy guarantees depend on the source: Arrow-backed inputs preserve buffers,
+while records built from Python lists or NumPy arrays will materialize once into
+Arrow buffers. The same applies to `StructScalar` inputs, which are normalized
+through Python objects before Arrow-mode export.
+For Parquet/Vortex sources, zero-copy also requires the on-disk struct schema
+to match `OME_ARROW_STRUCT`; non-strict schema normalization materializes via
+Python objects.
+
+## Optional dependencies
+
+CPU DLPack export uses Arrow buffers by default. For framework helpers and GPU
+paths, install only what you need:
+
+```bash
+pip install "ome-arrow[dlpack-torch]"  # torch only
+pip install "ome-arrow[dlpack-jax]"    # jax only
+pip install "ome-arrow[dlpack]"        # both
+```

docs/src/examples/learning_to_fly_with_ome-arrow.py

@@ -103,3 +103,41 @@
 oa_image = OMEArrow(data="../../../tests/data/idr0062A/6001240_labels.zarr")
 # show the image using pyvista
 oa_image.view(how="pyvista")
+
+# ## DLPack tensor export (advanced)
+# This is optional and requires torch: `pip install "ome-arrow[dlpack-torch]"`
+
+# +
+# examples of exporting OME-Arrow data into DLPack format for zero-copy
+import jax.numpy as jnp
+import torch
+
+oa = OMEArrow("example.ome.parquet")
+# -
+
+# %%time
+# DLPack Arrow mode: zero-copy 1D values buffer + reshape
+view = oa.tensor_view(t=0, z=0, c=0)
+cap = view.to_dlpack(mode="arrow")
+flat = torch.utils.dlpack.from_dlpack(cap)
+tensor = flat.reshape(view.shape)
+tensor.shape
+
+# %%time
+# DLPack NumPy mode: shaped tensor directly (still zero-copy when possible)
+# Layout quick reference:
+# - `C` = channels
+# - `H` = image height (Y axis)
+# - `W` = image width (X axis)
+view_chw = oa.tensor_view(t=0, z=0, layout="CHW")
+cap_chw = view_chw.to_dlpack(mode="numpy", contiguous=True)
+tensor_chw = torch.utils.dlpack.from_dlpack(cap_chw)
+tensor_chw.shape
+
+# %%time
+# DLPack Arrow mode: zero-copy 1D values buffer + reshape
+view = oa.tensor_view(t=0, z=0, c=0)
+caps = view.to_dlpack(mode="arrow")
+flat = jnp.from_dlpack(caps)
+arr = flat.reshape(view.shape)
+arr.shape

docs/src/index.md

@@ -13,4 +13,5 @@ caption: 'Contents:'
 maxdepth: 3
 ---
 python-api
+dlpack
 ```

docs/src/python-api.md

@@ -9,3 +9,12 @@ ome_arrow.meta
    :undoc-members:
    :show-inheritance:
 ```
+
+```{eval-rst}
+ome_arrow.tensor
+-------------------
+.. automodule:: src.ome_arrow.tensor
+   :members:
+   :undoc-members:
+   :show-inheritance:
+```

pyproject.toml

@@ -30,6 +30,16 @@ dependencies = [
   "pillow>=12",
   "pyarrow>=22",
 ]
+optional-dependencies.dlpack = [
+  "jax>=0.4",
+  "torch>=2.1",
+]
+optional-dependencies.dlpack-jax = [
+  "jax>=0.4",
+]
+optional-dependencies.dlpack-torch = [
+  "torch>=2.1",
+]
 optional-dependencies.viz = [
   "ipywidgets>=8.1.8",
   "jupyterlab-widgets>=3.0.16",
@@ -114,6 +124,8 @@ lint.per-file-ignores."*" = [ "ANN401", "C901", "PLC0415", "PLR0912", "PLR0913",
 lint.per-file-ignores."__init__.py" = [ "F401" ]
 # ignore docstring presence checks for docs
 lint.per-file-ignores."docs/*" = [ "D100", "D101", "D102", "D103", "D104", "D105", "D106", "D107" ]
+# allow notebook-style mid-file imports in paired examples
+lint.per-file-ignores."docs/src/examples/learning_to_fly_with_ome-arrow.py" = [ "E402" ]
 # ignore typing rules for tests
 lint.per-file-ignores."tests/*" = [ "ANN201", "E501", "PLR0913", "PLR2004" ]
 

src/ome_arrow/__init__.py

@@ -20,6 +20,7 @@
     to_ome_arrow,
 )
 from ome_arrow.meta import OME_ARROW_STRUCT, OME_ARROW_TAG_TYPE, OME_ARROW_TAG_VERSION
+from ome_arrow.tensor import TensorView
 from ome_arrow.utils import describe_ome_arrow, verify_ome_arrow
 from ome_arrow.view import view_matplotlib, view_pyvista