}
    device_to_host_stream =
        std::shared_ptr<se::Stream>(std::move(ptr_or_status.value()));
  }

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  xla_tensor->WaitForDefinitionEventOnStream(device_to_host_stream.get());

  // Transfer manager requires the shape of the shaped buffer to be the same as
  // literal shape except for the layout.  Set the literal to use xla_tensor's

            << " output buffers: " << xla_output->shaped_buffer().ToString();

    // Wait for definition event of the source tensor so the input buffers are
    // available.
    xla_input->WaitForDefinitionEventOnStream(dst_device_to_device_stream);

    // Wait for the destination tensor buffers to be ready, if they are not
    // available for an immediate write.

          << "Must have a stream available when using XLA tensors!";
      XlaTensor* xla_tensor = XlaTensor::FromTensor(t);
      CHECK(xla_tensor);
      xla_tensor->WaitForDefinitionEventOnStream(
          ctx->op_device_context()->stream());
    }

    arguments.emplace_back(device_shape, host_shape);
    xla::ExecutionInput& execution_input = arguments.back();