const xla::Shape& on_device_shape,
                                       xla::LocalClient* client,
                                       int device_ordinal) {
  xla::Shape on_host_shape =
      xla::ShapeUtil::DeviceShapeToHostShape(on_device_shape);
  xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape,
                                        client->backend().memory_allocator(),

        "XlaTensor is expected to have device memory allocated when "
        "computing padded shape");
  }

  const xla::Shape& on_device_shape =
      xla_tensor->shaped_buffer().on_device_shape();

  StatusHelper status;
  ApiConverter::StackHelper<XLA_Shape> se_shape(on_device_shape);
  ApiConverter::StackHelper<XLA_Shape> tpu_shape;
  stream_executor::tpu::ExecutorApiFn()->XlaShapeToTpuPaddedShapeFn(

  // shape. If a ShapedBuffer exists already (has_shaped_buffer() == true), it
  // is replaced and the managed memory deallocated.
  Status AllocateShapedBuffer(DataType dtype, const xla::Shape& on_device_shape,
                              xla::LocalClient* client, int device_ordinal);

  // Some Tensors can have complex on-device shapes, including tuple shapes. To

  // Computation output should always be a tuple.
  VLOG(2) << "Result tuple shape: " << output.on_host_shape().DebugString();
  VLOG(2) << "Result tuple shape (on device): "
          << output.on_device_shape().DebugString();
  CHECK_EQ(ctx->num_outputs(), compilation_result->outputs.size());

  // If the on-host-shape isn't a tuple, create a new single-element tuple

  if (xla_tensor == nullptr) {
    return TensorShapeToXLAShape(tensor.dtype(), tensor.shape(), shape);
  }

  const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer();
  *shape = shaped_buffer.on_device_shape();
  return absl::OkStatus();
}

// Caches a XlaDeviceAllocator per <backend, device ordinal> pair. A
// XlaDeviceAllocator is created on demand and is associated with a