🐛[BUG]: A pre-allocated Batch.empty buffer can only be filled once: every put() after the first is a silent no-op

[lil-lon]

Version

0.1.0 (also present on the latest main checkout)

On which installation method(s) does this occur?

Pip, Source

Describe the issue

A buffer created with Batch.empty(...) accepts exactly one put(). After the first successful put(), every subsequent put() silently copies nothing: num_graphs stops growing and no error or warning is raised. zero() is not required to trigger this: a plain put → put sequence drops the second copy too. The repro below uses put → zero → put only because it makes the symptom most visible (num_graphs stays 0).

Expected: a buffer with remaining capacity should accept further put()s, and zero() should restore the freshly-allocated state. At minimum, a dropped copy should raise or warn rather than silently no-op.

This also affects in-tree consumers that reuse a Batch.empty buffer:

• GPUBuffer uses the same pattern but is not affected in practice: it already carries an explicit workaround (_restore_batch_ptr_capacity) that re-extends the trimmed batch_ptr after every write. This workaround was introduced in Buffer sync semantics #4, so the trimming behavior itself seems to be known; fixing it in put would make the workaround removable.
• The distributed pipeline: _CommunicationMixin reuses pre-allocated send/recv buffers across steps with this exact pattern (_recv_to_batch does put then zero; _batch_to_buffer does repeated put), so every transfer after the first one would be silently dropped.

Minimum reproducible example

import torch
from ase.build import molecule

from nvalchemi.data import AtomicData, Batch


def atoms_to_data(atoms) -> AtomicData:
    data = AtomicData.from_atoms(atoms)
    data.forces = torch.zeros(data.num_nodes, 3)
    data.energy = torch.zeros(1, 1)
    return data


def put_one(buf: Batch, name: str) -> int:
    src = Batch.from_data_list([atoms_to_data(molecule(name))])
    buf.put(src, mask=torch.ones(1, dtype=torch.bool))
    return buf.num_graphs


template = Batch.from_data_list([atoms_to_data(molecule("H2O"))])
buf = Batch.empty(num_systems=4, num_nodes=64, num_edges=0,
                  template=template, device="cpu")

print("put #1 -> num_graphs =", put_one(buf, "H2O"))

buf.zero()
print("zero() -> num_graphs =", buf.num_graphs)

print("put #2 -> num_graphs =", put_one(buf, "CH4"), " (expected 1)")
print("put #3 -> num_graphs =", put_one(buf, "NH3"), " (expected 2)")

Relevant log output

put #1 -> num_graphs = 1
zero() -> num_graphs = 0
put #2 -> num_graphs = 0  (expected 1)
put #3 -> num_graphs = 0  (expected 2)

Environment details

Device-independent: reproduces on both CPU and GPU.
Not tied to any specific CUDA / GPU / PyTorch / Warp version.

Additional context

I have already tracked down the root cause (the first successful put() truncates the pre-allocated batch_ptr in SegmentedLevelStorage.put, so all later capacity checks fail) and have a fix in hand. I understand direct code contributions are not accepted during the public beta. Happy to open a PR once they are, or earlier if a maintainer assigns this issue to me.