Add Tensor.copyDataInto to Java API

[jgibson2]

Summary

Adds Tensor.copyDataInto(... dst) to the Android Java API for the float32 and float16 dtypes. It copies the tensor's data into a caller-provided destination buffer instead of allocating a fresh float[] per call (as getDataAsFloatArray() does today). The same pattern is repeated for other types.

Motivation

While profiling depth inference on Android with Perfetto, output extraction was a meaningful contributor to ART GC pressure. Each call to output.toTensor().dataAsFloatArray allocates a new Java float[] sized to the tensor's element count and bulk-copies from the underlying off-heap buffer into it.

The native side already exposes the underlying FloatBuffer directly (zero-copy view of the C++ tensor's data_ptr()), so the only thing missing was a public way for callers to drain it into a destination buffer they already own and reuse across calls.

API

public void copyDataInto(FloatBuffer dst)

• Implemented on all datatypes

Caller-side usage example

// One-time setup
FloatBuffer depthBuf = Tensor.allocateFloatBuffer(numelDepth);

// Per inference
EValue[] outputs = module.forward(...);
depthBuf.rewind();
outputs[0].toTensor().copyDataInto(depthBuf);   // no allocation
// ... read from depthBuf ...

Test plan

• Added unit tests in TensorTest.kt:
  • testCopyDataIntoFloat32 — round-trip with reuse across two calls
  • testCopyDataIntoFloat32_writesAtDstPosition — verifies the call writes at dst.position() and advances it (does not overwrite from index 0)
  • testCopyDataIntoFloat32_overflow — BufferOverflowException on undersized destination
  • testCopyDataIntoFloat16 — verifies fp16→fp32 widening matches getDataAsFloatArray
  • testCopyDataIntoFloat_unsupportedDtype — IllegalStateException from base default for non-float dtypes

This PR was authored with Claude.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19171

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 1 New Failure, 1 Cancelled Job, 1 Unrelated Failure

As of commit 068212e with merge base 79b799d ():

NEW FAILURE - The following job has failed:

• pull / unittest / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_linear_model

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-models-linux (phi_4_mini, portable, linux.4xlarge.memory) / linux-job (gh)
  ##[error]The operation was canceled.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[kirklandsign]

@claude review

[claude]

Claude finished @kirklandsign's task in 1m 17s —— View job

---

Code Review: Add Tensor.copyDataInto to Java API

• Read changed files and gather context
• Review Tensor.java changes
• Review TensorTest.kt changes
• Post review feedback

---

Overall: This is a clean, well-motivated PR. The API design is consistent with existing patterns, and test coverage is thorough. A few observations below, mostly minor.

What's good

• Motivation is solid — eliminating per-frame float[] allocations in inference loops is a real win for GC pressure on Android.
• API mirrors existing patterns — the copyDataInto overrides follow the same data.rewind(); dst.put(data) approach as getDataAs*Array(), just with a caller-provided destination.
• int8/uint8 disambiguation via copyDataIntoUnsigned is the right call since both share ByteBuffer.
• Test coverage is comprehensive — happy path, reuse, position-awareness, overflow, cross-dtype error messages, and signed/unsigned discrimination are all covered.
• Float16 → FloatBuffer widening correctly mirrors the existing getDataAsFloatArray() logic for Tensor_float16.

Issues

1. Thread safety of data.rewind() (pre-existing, not a blocker)
   Tensor.java:829 — Every copyDataInto override calls data.rewind() on the tensor's internal buffer before reading. This is the same pattern getDataAs*Array() already uses, so it's consistent. But worth knowing: java.nio.Buffer operations are not thread-safe. If two threads call copyDataInto (or any mix of copyDataInto/getDataAs*Array) on the same Tensor concurrently, they can corrupt each other's reads. Since the PR description mentions "steady-state inference loops" — if callers ever share a Tensor across threads, they'll need external synchronization. This is a pre-existing concern, not introduced here.

2. Partial write on overflow for float16 copyDataInto(FloatBuffer)
Tensor.java:893-898 — The element-by-element loop for fp16→fp32 widening means a BufferOverflowException can leave dst partially written and data's position somewhere in the middle. The bulk dst.put(data) used by other dtypes is atomic at the buffer level (it checks remaining capacity upfront before transferring). This is a minor inconsistency — in practice callers should size their buffers correctly — but you could guard with an upfront check:

   if (dst.remaining() < data.remaining()) {
       throw new BufferOverflowException();
   }

   Not a blocker, but would make the fp16 path behave identically to other dtypes on undersized buffers.

Nits

• The Javadoc on the FloatBuffer base method (lines 557–576) is quite detailed compared to the other copyDataInto overloads. This makes sense since it's the one most callers will reach for, just noting the asymmetry is intentional.
• testCopyDataIntoFloat64 and testCopyDataIntoFloat16RawBits manually construct DoubleBuffer/ShortBuffer via ByteBuffer.allocateDirect(...).order(...).asFooBuffer() — Tensor.allocateDoubleBuffer and Tensor.allocateHalfBuffer already exist and would shorten these.

Summary

Looks good to land. The only suggestion worth considering before merge is the upfront capacity check in Tensor_float16.copyDataInto(FloatBuffer) to avoid partial writes (issue #2), but it's not a must-have.

[jgibson2]

@pytorchbot label "release notes: none"

[jgibson2]

@claude re-review

[mergennachin]

@claude Review this PR

[claude]

Claude finished @mergennachin's task in 1m 1s —— View job

---

Code Review: Add Tensor.copyDataInto to Java API

• Read changed files and gather context
• Review Tensor.java changes
• Review TensorTest.kt changes
• Post review feedback

---

Overall: This is a well-designed, well-tested addition. The API is consistent with existing patterns, the motivation (reducing GC pressure in inference loops) is sound, and the test coverage is thorough. Ready to land with one minor suggestion.

What's good

• Complete dtype coverage — every supported tensor subclass (float32, float16, int8, uint8, int32, int64, float64) gets a copyDataInto override, plus the float16 → FloatBuffer widening path. No gaps.
• copyDataIntoUnsigned disambiguation — since both int8 and uint8 back onto ByteBuffer, giving unsigned its own method name avoids ambiguity and makes misuse a compile-time error rather than a runtime surprise. The tests at lines TensorTest.kt:453-459 and TensorTest.kt:475-481 verify cross-type rejection, which is great.
• Upfront capacity check for fp16 widening (Tensor.java:899-901) — the previous review flagged the partial-write risk and it's been addressed. The testCopyDataIntoFloat16_overflowIsAtomic test (TensorTest.kt:350-362) verifies dst.position() stays at 0 on failure. Well done.
• Test coverage is comprehensive — happy path, buffer reuse, position-awareness, overflow, cross-dtype rejection, and signed/unsigned discrimination are all covered across 15 new test methods.

Minor suggestion

data.rewind() could be data.position(0) (or left as-is) — Buffer.rewind() clears the mark, while position(0) preserves it. In practice no code in this class uses marks, so this is purely a hygiene note and consistent with the existing getDataAs*Array() methods. Not worth changing — just noting for awareness.

Summary

Clean PR. The API mirrors the existing getDataAs*Array() pattern, the fp16 widening path has all-or-nothing overflow semantics, and the int8/uint8 disambiguation is well handled. Looks good to land.