Explore DLPack ownership semantics

[kylebarron]

d-v-b#1

[kylebarron]

Spec written by Claude in discussion with me:

---

DLPack ownership semantics for Tensor

Issue: #108 · External feedback: d-v-b#1
Date: 2026-07-23 · Status: Design approved, not yet implemented

Problem

PyTensor holds decoded chunk data as bytes: Bytes — a refcounted, shared,
immutable buffer — and hands out several read-only, zero-copy faces of that
same allocation: buffer(), to_numpy()/__array__, and the PEP 3118
buffer-protocol export (__getbuffer__, hard-coded readonly=1). The DLPack
face (__dlpack__) currently returns a legacy DLManagedTensor, which has
no read-only concept, so np.from_dlpack(tensor) is writable-by-omission today.

d-v-b#1 asks us to make the DLPack export deliberately writable and
zero-copy, so the zarr-python engine binding's read path can be
np.from_dlpack(result) without its writability copy-guard firing. zarr-python's
API contract is that reads return writable arrays.

The load-bearing constraint

Zero-copy + writable requires unique ownership. There is no sound
(zero-copy + writable + shared).

PyTensor is Clone and exposes multiple read-only aliases of the same Bytes.
Exporting a writable alias while any of those can coexist means writes leak
through faces that advertised immutability. It is memory-safe (all bytes are
valid u8) but it violates the immutability contract those faces promise — the
kind of invisible convention the project's "parse, don't validate" philosophy
exists to design away.

The escape is forced: a sound writable export must move the unique buffer out
of the PyTensor, leaving the tensor tombstoned. is_unique()-checking and
moving-out are two sides of the same requirement.

DLPack background (the parts that matter here)

1. Read-only is a versioned-only capability. DLManagedTensor (legacy) has
   no flags field — a legacy capsule is writable-by-convention.
   DLManagedTensorVersioned adds flags, where READ_ONLY and IS_COPIED
   live. Controlling mutability ⟺ exporting versioned.

2. __dlpack__ is a negotiation. The modern protocol is
   __dlpack__(*, stream, max_version, dl_device, copy). max_version tells us
   whether the consumer accepts versioned; dl_device can request another device
   (we error — CPU only); copy is the consumer's only mutability/aliasing lever:
   True = "I intend to own/mutate, give me a copy", False = "must be
   zero-copy or fail", None = "producer's choice". There is no consumer
   channel to request a writable alias — READ_ONLY flows producer → consumer
   only. That absence is why an explicit, producer-side "take ownership" door is
   the right shape.

3. Keep-alive = manager_ctx + deleter. Whatever is moved into the manager
   context stays alive until the consumer runs the deleter — on the consumer's
   thread, possibly long after the PyTensor is gone.

4. DLPack is natively one-shot. A consumed capsule is renamed
   used_dltensor; a managed tensor is meant to be taken once. An explicit,
   consuming move door is with the grain of the protocol.

Design: two explicit doors

Door 1 — Tensor.__dlpack__(): repeatable, immutable

Rust exports stay immutable by default. Behavior is driven entirely by the
standard copy kwarg, exporting versioned when max_version allows (legacy
fallback otherwise):

consumer call	result	flag	buffer
copy=None/False	read-only, shared, zero-copy	READ_ONLY	Bytes clone (kept)
copy=True	writable, copied, owned	IS_COPIED	fresh unique copy

• Non-consuming and repeatable — consistent with buffer(), to_numpy(), and
  the buffer-protocol face.
• copy=False is satisfiable via the shared read-only view (zero-copy), so it
  does not raise.
• dl_device other than CPU → BufferError.

Door 2 — Tensor.into_dlpack(): explicit one-shot move

The only path to writable zero-copy. Ownership transfer is explicit.

• Gate on Bytes::is_unique(). If a buffer-protocol view or a Tensor clone is
  still outstanding, the buffer is not unique → raise (e.g. BufferError /
  ValueError: "Tensor buffer is shared; cannot move out for a writable export.
  Use __dlpack__(copy=True) for a writable copy.").
• If unique: take the Bytes out of the tensor's slot (tombstone — see
  below) and move it into a small intermediate handle, OwnedTensor, that owns
  the Bytes and implements __dlpack__ / __dlpack_device__.
• OwnedTensor.__dlpack__ yields a writable, zero-copy, versioned capsule
  (no READ_ONLY) and moves the Bytes into the capsule's manager context. It
  is itself one-shot.

The intermediate OwnedTensor exists so that np.from_dlpack(t.into_dlpack())
can still negotiate max_version/copy on the writable path — into_dlpack()
returns an object with __dlpack__, not a pre-built capsule.

Uniqueness is checked at the into_dlpack() call (the take-ownership moment). On
failure the tensor's slot is left intact (not tombstoned).

Buffer representation

Hold Bytes, not BytesMut. Rust never mutates the buffer; mutation happens
consumer-side after the move, sound purely because is_unique() guaranteed
refcount 1 at move time, so no Rust-side alias remains. Bytes with count 1
dropping in the deleter frees the allocation. BytesMut would falsely imply Rust
intends to write.

To make the buffer takeable while keeping the read-only faces working, the slot
becomes interior-mutable: Mutex<Option<Bytes>> (the GIL already serializes
access from Python; the Mutex only satisfies the frozen pyclass Sync
requirement — the deleter never touches this slot, it owns its own moved-out
Bytes). Metadata (shape, dtype) stays plain and survives consumption.

Faces after consumption (tombstone)

• Read-only faces (buffer(), to_numpy(), __array__, __getbuffer__) and
  __dlpack__ clone the Bytes out of the Option. If the slot is None
  (already moved out by into_dlpack()), they raise "Tensor already consumed".
• Metadata getters (shape, dtype) continue to work.

Relationship to d-v-b#1

This design deliberately does not make the bare np.from_dlpack(result)
writable. Under Door 1 the plain call (copy=None) is read-only, so the binding
must opt in explicitly:

• np.from_dlpack(result.into_dlpack()) — writable, zero-copy (the move), or
• np.from_dlpack(result, copy=True) — writable, but a copy (their current
  fallback, now negotiated honestly).

This is intentional pushback: Rust exports are immutable by default; taking
ownership is explicit. In the binding's read path result is freshly decoded
and unique, so into_dlpack() succeeds and gives true zero-copy.

Sequencing

• PR-1 (finish chore: bump dlpark to 0.8 #102, branch kyle/bump-dlpark): purely mechanical bump to
  dlpark 0.8, staying on legacy Dlpack. No behavior change, no versioned,
  no writability. Get onto 0.8 and green.
• PR-2 (this issue): on top of PR-1, migrate __dlpack__ to versioned +
  copy negotiation (Door 1) and add into_dlpack()/OwnedTensor (Door 2),
  plus the interior-mutable slot and tombstoning.

To verify at implementation time

• Exact dlpark 0.8 Builder API for producing a versioned managed tensor
  and setting/clearing DlpackFlags (READ_ONLY, IS_COPIED). The 0.8 rustdoc
  confirms builder/legacy/versioned modules and a DlpackFlags struct;
  method signatures were not captured here.
• dlpark 0.8 Default for Flags — the 0.7 fork defaulted to READ_ONLY, a
  footgun to avoid inheriting silently. Set flags explicitly on every export.
• numpy's forwarding of from_dlpack(x, copy=..., device=...) into
  x.__dlpack__(copy=..., dl_device=...) (numpy ≥ 2.1) — confirm the version
  floor for the copy=True path.
• Capsule naming / one-shot rename (dltensor → used_dltensor,
  dltensor_versioned) is handled by dlpark's IntoPyObject, not us.