namespace TF {

// Returns whether type can be further refined.
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);

def ReorderReshapeDequantQuant : Pat<(TF_ReshapeOp:$old_value
              (TFL_DequantizeOp (TFL_QuantizeOp $input, $qtype)), $shape),
          (TFL_DequantizeOp
              (TFL_QuantizeOp (TF_ReshapeOp $input, $shape),
              (UpdateShapeWithAxis<-1> $qtype, $old_value))),
          [(CanUpdateShapeWithAxis<-1> $qtype, $old_value)]>;

// CHECK: stablehlo.return %[[MAX]]

// Check that the attributes window_dimensions & window_strides are also
// permutated to match the new input shape.
// CHECK: (tensor<1x16x16x4xf32>, tensor<f32>) -> tensor<1x8x8x4xf32>

// Check that a `stablehlo.transpose` is added to the result to match the shape
// of the users.

  for (const auto& shape : llvm::enumerate(input_shapes_vector)) {
    if (!shape.value().has_value()) {
      TF_RETURN_IF_ERROR(TensorShapeUtils::MakeShape(
          static_cast<int*>(nullptr), 0, &arg_shapes[shape.index()].shape));
      continue;
    }
    TF_RETURN_IF_ERROR(TensorShapeUtils::MakeShape(
        *shape.value(), &arg_shapes[shape.index()].shape));
  }

  return absl::OkStatus();
}

  // CHECK-DAG:  %[[CSTR_LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:  %[[CSTR_RHS_SHAPE:.+]] = shape.shape_of %arg1
  // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[CSTR_LHS_SHAPE]], %[[CSTR_RHS_SHAPE]]
  // CHECK-NEXT: shape.assuming %[[WITNESS:.+]]
  // CHECK-DAG:    %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:    %[[RHS_SHAPE:.+]] = shape.shape_of %arg1

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

  if (shape.rank() == 0 ||
      (shape.dimensions_size() == 1 && shape.dimensions(0) == 1)) {
    dim_sizes = "[1]";
    indices = "[0]";
  } else {
    for (int dim = 0; dim < shape.dimensions_size(); ++dim) {
      dim_vars.push_back(absl::StrCat("size_t dim", dim));
      dim_sizes += absl::StrCat("[", shape.dimensions(dim), "]");
      indices += absl::StrCat("[dim", dim, "]");
      count *= shape.dimensions(dim);

}

// Create an empty tensor of type 'dtype'. 'shape' can be arbitrary, but has to
// result in a zero-sized tensor.
static TF_Tensor* EmptyTensor(TF_DataType dtype,
                              const tensorflow::TensorShape& shape) {
  static char empty;
  int64_t nelems = 1;
  std::vector<int64_t> dims;
  auto shape_dims = shape.dims();
  dims.reserve(shape_dims);
  for (int i = 0; i < shape_dims; ++i) {

      auto shape = llvm::to_vector<4>(inputType.getShape());
      llvm::SmallVector<int64_t, 4> bounds(shape.begin(), shape.end());
      // Mark the dim as dynamic dim.
      shape[0] = ShapedType::kDynamic;
      auto extensions =
          mhlo::TypeExtensionsAttr::get(func->getContext(), bounds);
      auto resultType =
          RankedTensorType::get(shape, inputType.getElementType(), extensions);

  %cst = "tf.Const"() {value = dense<3.000000e-01> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
  %0 = "tf.XlaCallModule"(%arg0, %cst) {
    Sout = [#tf_type.shape<1x2x2x2>], _entry_function = @composite_conv_fn,
    _original_entry_function = "composite_conv_fn",
    _stablehlo_module_attrs = {}, _quantization_method = "weight_only_ptq { }",