Assemble loader tensors directly into their spill destination

[mkolodner-sc]

What

TFRecordDataLoader.load_as_torch_tensors assembles each entity's tensor by accumulating per-batch tensors in a list and joining them with tf.concat. The final tf.concat allocates the whole result while the list of batches is still alive, so the assembled tensor momentarily coexists with a second full-size copy of itself in anonymous memory.

When tensor spilling is enabled, this change allocates the full-size destination up front via allocate_disk_backed(...) and streams each batch into its slice, freeing the batch as it goes — so the result never coexists with a second full-size copy. When a disk-backed buffer is not used (spilling disabled, below the size threshold, or an unmapped dtype), it falls back to the existing tf.concat path unchanged, so behaviour is identical for callers that do not spill.

This is the "buffer that does not exist yet" case allocate_disk_backed already documents ("allocating the destination as the file from the start means those bytes never occupy anonymous memory at all") — the loader just wasn't using it for the final assembly yet.

Why

The assembly step is a peak-memory hotspot at load time. Measured on a 7.6 GiB fp32 shard (process RSS high-water-mark across the assembly step):

path	peak RSS
tf.concat (before)	16.3 GiB
assemble into destination (this PR)	8.7 GiB

The transient second copy is removed — peak drops from ~2× the shard to ~1×.

Tests

Adds test_load_as_torch_tensors_assembles_into_spilled_destination: with spilling enabled the loaded ids / features / labels are byte-identical to the in-memory (tf.concat) result and are disk-backed; with spilling disabled, behaviour is unchanged. The existing dataloaders_test.py cases continue to cover the non-spilling (fallback) path.

[kmontemayor2-sc]

Thanks Matt! LGTM provided the comments are addressed :)