xla::Shape subshape =
        xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first);
    uint64 size =
        client->backend().transfer_manager()->GetByteSizeRequirement(subshape);
    TF_ASSIGN_OR_RETURN(se::OwningDeviceMemory buffer,
                        client->backend().memory_allocator()->Allocate(
                            device_ordinal, size, /*retry_on_failure=*/false,

    // Copy host -> device. (Empty tensors don't have backing buffers.)
    // Manually allocate memory so we can allocate as much memory as the device
    // requires (as given by GetByteSizeRequirement). This avoids
    // XlaTransferManager having to reallocate the device buffer later if
    // XlaTransferManager is used.
    VLOG(1) << "Constant output tensor on device";

    TF_RETURN_IF_ERROR(