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.

                                  PatternRewriter &rewriter) {
  // Fail for non-static shapes
  if (!pad_op.getOperand().getType().hasStaticShape() ||
      !pad_op.getResult().getType().hasStaticShape() ||
      !pad_op.getPaddingValue().getType().hasStaticShape()) {
    return rewriter.notifyMatchFailure(pad_op, "dynamic shapes not supported");
  }

  // Check if the operand is also a Pad op
  auto parent_pad =

      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.

  // Fuse select(broadcast_to(input, shape), x, y) -> selectV2(input, x, y)
  // Also, fuse selectv2(broadcast_to(input, shape), x, y) -> selectV2(input, x, y)
  // It is safe to perform this transform here because-
  // the shapes of `pre_broadcast` and `dim` must be broadcast
  // compatible for the `broadcast_to` op to be valid.
  // And considering, `shape(post_broadcast)` == `shape(%input1)`,

		resources []config.ResourceSpec
		shapes    []config.UtilizationShapePoint
		wantErrs  field.ErrorList
	}{
		{
			name:   "no shapes",
			shapes: nil,
			wantErrs: field.ErrorList{
				{
					Type:  field.ErrorTypeRequired,
					Field: "scoringStrategy.shape",
				},
			},
		},
		{
			name:   "weight greater than max",
			shapes: defaultShape,
			resources: []config.ResourceSpec{

    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

  // Check all elements besides at concat_dim match across all shape tensors.
  SmallVector<int32_t, 4> shape0;
  shape0.reserve(num_dims);
  for (int32_t dim : shapes.front().getValues<int32_t>()) shape0.push_back(dim);

  for (DenseIntElementsAttr shape : llvm::drop_begin(shapes, 1)) {
    for (const auto& dims_and_idx : llvm::enumerate(llvm::zip(shape0, shape))) {
      if (dims_and_idx.index() == concat_dim) continue;

// 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.

namespace tensorflow {

class XlaDevice : public LocalDevice {
 public:
  // Given a tensor, sets `xla::Shape*` the shape of tensor's representation
  // on device, fully padded. On error, the contents of `xla::Shape*`
  // are undefined.
  typedef std::function<Status(const Tensor&, xla::Shape*)> PaddedShapeFn;

  // Wrapper class to store metadata about the XlaDevice, where it can be