//
  // For example,
  //
  //   %shape  = tf.Shape(%arg)                   // %arg: tensor<?x2x3x1xf32>
  //   %height = tf.StridedSlice(%shape, 1, 2, 1)
  //
  // In this case %height can be replaced with a constant 2.
  //
  // Or,
  //
  //   %shape  = tf.Shape(%arg)                   // %arg: tensor<?x2x3x1xf32>
  //   %spatial_shape = tf.StridedSlice(%shape, 1, 3, 1)
  //

    // Merge node causes a loop so we remove NextIteration->Merge edge before
    // performing shape inference. But removing those edges also prevents us
    // from inferring output shape for Merge node (we need shapes for all its
    // inputs).
    // For loop invariant resource input's Merge node, we set output resource
    // shape as Enter node's resource shape.
    // TODO(b/129367850): clean this up.

// RUN: tf-opt -split-input-file -verify-diagnostics --tf-shape-inference %s | FileCheck %s

module attributes {tf.versions = {producer = 888 : i32}} {
// CHECK-LABEL: testConv2dShapeValidPadding
func.func @testConv2dShapeValidPadding(%arg0: tensor<1x112x80x128xf32>, %arg1: tensor<128x3x3x128xf32>, %arg2: tensor<128xf32>) -> tensor<1x?x?x128xf32> {

// Extracts shape from XlaArgument as TensorShape. If shape is a xla::Shape,
// that is converted to a TensorShape.
absl::StatusOr<TensorShape> GetTensorShapeFromXlaArgument(
    const XlaArgument& arg) {
  if (absl::holds_alternative<xla::Shape>(arg.shape)) {
    TensorShape arg_shape;
    TF_RETURN_IF_ERROR(
        XLAShapeToTensorShape(std::get<xla::Shape>(arg.shape), &arg_shape));
    return arg_shape;

^bb0(%arg0: tensor<4x4x4xf32>) :
  %shape0 = arith.constant dense<[16, 4]> : tensor<2xi32>
  %shape1 = arith.constant dense<[64]> : tensor<1xi32>
  %0 = "tfl.reshape"(%arg0, %shape0) : (tensor<4x4x4xf32>, tensor<2xi32>) -> tensor<16x4xf32>
  %1 = "tfl.reshape"(%0, %shape1) : (tensor<16x4xf32>, tensor<1xi32>) -> tensor<64xf32>
  %2 = "tfl.reshape"(%0, %shape1) : (tensor<16x4xf32>, tensor<1xi32>) -> tensor<64xf32>

  }

  std::vector<std::optional<std::vector<int>>> shapes;
  TF_RETURN_IF_ERROR(::tensorflow::ParseNodeShapes(input_shapes, shapes));

  for (const auto& shape : shapes) {
    if (!shape) {
      return absl::AbortedError("Missing input argument shapes");
    }
    parsed_shapes.push_back(SmallVector<int64_t>(shape->begin(), shape->end()));
  }
  return parsed_shapes;
}

// 'shape' is the original shape with padding to match result shape.
int64_t GetElementIndex(const std::vector<int64_t>& shape,
                        const std::vector<int64_t>& current_index) {
  int64_t ind = 0;
  int64_t mul = 1;
  for (int i = shape.size() - 1; i >= 0; --i) {
    ind += (current_index[i] % shape[i]) * mul;
    mul *= shape[i];
  }
  return ind;
}

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.