bool CanBeRefined(Type type);

// Returns a new arg type based on the shape and element type. If there are
// dynamic bounds attribute to the arg, update the bounds based on the shape
// as well.
Type GetNewArgType(Type old_arg_type, ArrayRef<int64_t> shape,
                   Type element_type, mlir::MLIRContext* context);

// Refines all the shapes in a module, skipping the inference for all ops
// whose type is in ops_to_skip.

      se::Stream* stream, se::DeviceMemoryBase* device_memory_base,
      const xla::Shape& shape,
      tsl::AsyncValueRef<std::unique_ptr<se::Event>>& done_event)
      : stream_(stream),
        device_memory_base_(device_memory_base),
        shape_(shape),
        done_event_(done_event) {}

  // Copies 'cpu_tensor' to `device_memory_base_` with `shape_`.
  // `device_tensor` is unused.

    TF_RET_CHECK(sit != shape_info.end())
        << "Missing shape information for node " << node->name();
    std::vector<PartialTensorShape> shapes;
    for (const auto& output : sit->second) shapes.push_back(output.shape);

    auto it = expected_shapes.find(node->name());
    if (it != expected_shapes.end()) {
      if (!PartialTensorShapeUtils::AreIdentical(shapes, it->second)) {
        return errors::InvalidArgument(

namespace tensorflow {

struct InferredShape {
  // Shape of the argument tensor.
  PartialTensorShape shape;

  // If the argument is a resource variable, the type and shape of the
  // variable's value.
  DataType handle_type = DT_INVALID;
  PartialTensorShape handle_shape;
};
typedef std::unordered_map<string, std::vector<InferredShape>> GraphShapeInfo;

// Infer shapes for all Tensors in a graph, and save them in a map.  The vector

// Converts a shape from MLIR to a TensorFlow tensor shape proto.
void ConvertToTensorShapeProto(llvm::ArrayRef<int64_t> shape,
                               TensorShapeProto* output_shape);

// Converts an MLIR type to a TensorFlow tensor shape.
PartialTensorShape ConvertTypeToTensorShape(const mlir::Type& type);

// Converts an MLIR shaped type to a TensorFlow shape attribute.

 public:
  XlaHostRecvDeviceContext(
      se::Stream* stream, const se::DeviceMemoryBase& device_memory_base,
      const xla::Shape& shape,
      tsl::AsyncValueRef<std::unique_ptr<se::Event>>& done_event)
      : stream_(stream),
        device_memory_base_(device_memory_base),
        shape_(shape),
        done_event_(done_event) {}

  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,

#include "xla/client/local_client.h"

namespace tensorflow {
namespace {
std::vector<const xla::Shape*> GetShapePointers(
    absl::Span<const xla::Shape> shapes) {
  std::vector<const xla::Shape*> shape_ptrs;
  shape_ptrs.reserve(shapes.size());
  for (const auto& shape : shapes) {
    shape_ptrs.push_back(&shape);
  }
  return shape_ptrs;
}
}  // namespace

absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>

// RUN: tf-opt %s -tf-shape-inference=input-arg-shapes=1 -verify-diagnostics -split-input-file | FileCheck %s
// RUN: not tf-opt %s -tf-shape-inference=input-arg-shapes=* 2>&1 | FileCheck --check-prefix=INPUT_ARG_SHAPES_ERROR %s
// INPUT_ARG_SHAPES_ERROR: Missing input argument shapes

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    // CHECK-LABEL: func.func @main

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2
        }
        dim {

  // CHECK-SAME: key = ""
  // CHECK-SAME: recv_key = "host_compute_channel_recv"
  // CHECK-SAME: send_key = "host_compute_channel_send"
  // CHECK-SAME: shape_inference_graph = @host_func
  // CHECK-SAME: shapes = [#tf_type.shape<*>, #tf_type.shape<3x?>]
  // CHECK-SAME: tpu_core = 0 : i64
  // CHECK: func @host_func
  // CHECK: %[[RECV_OUTPUT:[0-9]*]]:2 = "tf._XlaRecvAtHost"
  // CHECK-SAME: key = "host_compute_channel_send"