bool use_multiple_streams, std::shared_ptr<se::Event> definition_event) {
  xla::ShapeIndex output_index = input_output_alias.shape().IsTuple()
                                     ? xla::ShapeIndex({output_num})
                                     : xla::ShapeIndex({});
  CHECK(input_output_alias.shape().IsTuple() || output_num == 0);
  if (std::optional<xla::HloInputOutputAliasConfig::Alias> alias =

      // If the representation is a tuple, we also must wait for the tuple index
      // buffers to be available on the destination host to device transfer
      // stream.
      if (xla_output->shaped_buffer().on_device_shape().IsTuple()) {
        TF_RETURN_IF_ERROR(dst_xla_context->host_to_device_stream()->WaitFor(
            dst_compute_stream));
      }
    }

    for (const auto& leaf : xla_input->shaped_buffer().buffers().leaves()) {

      compilation_result->computation->GetProgramShape();
  TF_RET_CHECK(program_shape_or_status.ok());

  xla::ProgramShape& program_shape = program_shape_or_status.value();
  if (!program_shape.result().IsTuple())
    return errors::Internal("Expect result to have tuple shape");

  xla::HloInputOutputAliasConfig config(program_shape.result());
  for (auto alias : output_to_input_alias) {
    if (use_tuple_args) {