#5581·burn

[burn 0.21 / burn-fusion] Feature request: public API to construct a FusionTensor from a base-backend tensor (Client::adopt_base_tensor)

Author: Harry-ZhouCreated Sep 6, 2026Updated Sep 17, 2026

Environment

  • burn 0.21.0 (burn-fusion 0.21.0, burn-cubecl 0.21.0, burn-cuda 0.21.0)
  • Any Fusion<B> backend; we hit this with burn_cuda::Cuda (which aliases Fusion<CubeBackend<CudaRuntime, …>> in default builds)

Summary

The fusion → base direction is already publicly solvable (thank you — this works well):

rust
// burn_fusion::client::Client::resolve_tensor_float — pub
let base: B::FloatTensorPrimitive = fusion_tensor.client
    .resolve_tensor_float::<B>(fusion_tensor);  // device-side, drains the stream

The base → fusion direction has no public path:

  • FusionTensor::new(id, shape, dtype, client, stream) is pub(crate);
  • the count: Arc<AtomicU32> field is private, so the struct cannot even be built via literal syntax from outside the crate;
  • we found no public API that registers an existing base tensor handle into the fusion pool.

What we need

One of:

  1. Client::adopt_base_tensor(handle: B::FloatTensorPrimitive) -> FusionTensor<R> — the mirror of resolve_tensor_float (attach the base primitive to the fusion client, fresh id, current stream), or
  2. making FusionTensor::new pub (plus a public constructor for count), or
  3. a Tensor<Fusion<B>, D>::from_base(Tensor<B, D>) convenience on the burn side.

Use case (why the reverse direction matters)

We ship custom #[cube(launch_unchecked)] kernels that operate on raw ComputeClient + Handle (obtained from the unfused CubeBackend's CubeTensor, whose client/handle fields are public — thank you). For models running on the default fusion-wrapped backend:

  • inputs are now solved device-side via resolve_tensor_float (zero host round-trip);
  • outputs of a bridge kernel live on the unfused backend. Returning them to the fusion model currently requires a host round-trip (into_data()from_data()), because there is no public way to wrap a base tensor back into a FusionTensor.

With adopt_base_tensor, a bridge kernel's result would re-enter the fusion graph device-side and the whole pattern becomes allocation-free in both directions.

Minimal demonstration of the asymmetry

rust
use burn::tensor::{backend::Backend, Tensor};

type Unfused = CubeBackend<CudaRuntime, f32, i32, u8>;
type Fused = burn_cuda::Cuda; // Fusion<Unfused> in default builds

// fusion → base: public, device-side ✅
let base_t: Tensor<Unfused, 2> = fusion_to_cube(&fusion_t);

// base → fusion: no public path ❌
// let fusion_t: Tensor<Fused, 2> = Tensor::from_primitive(
//     FusionTensor::new(...)  // pub(crate) — cannot construct
// );

Workaround we ship today

Host round-trip for outputs (Tensor::into_data() on the unfused tensor, Tensor::from_data() on the fusion backend). Correct but costs a device→host→device copy per bridge call.

Notes

  • We deliberately do not ask to pierce the fusion pool internals — only to adopt an already-materialized base primitive under a fresh id.
  • If a public path already exists that we missed, a pointer would fully resolve this request.