tensor_spec.TensorSpec(
                shape=shape, dtype=dtypes.float32, name='input_tensor'
            )
        ),
    )
    return model

  # Prepares sample einsum input data shapes.
  # This function returns:
  # 1. Shape for input 1
  # 2. Shape for input 2
  # 3. Shape for bias
  # 4. Signature for input 1 (Could contain None dimension)

    bool lower_to_xla_hlo = true, bool allow_partial_conversion = false);

// Helper struct representing argument tensor or resource handle shapes.
struct TensorOrResourceShape {
  TensorShape shape;
  bool is_resource = false;
};

// Refine MLIR types based on new shape information.
ABSL_DEPRECATED("Not meant to be used directly and should be a util.")
Status RefineShapes(llvm::ArrayRef<TensorOrResourceShape> arg_shapes,

# TODO(fengliuai): Figure out what to do here. A "golden" end-to-end test is overly fragile.

  return transposed_type;
}

// Returns shape of a ranked tensor.
// Precondition: output_val's is ranked tensor.
// Returns a truncated shape when `truncate` is set to true.
inline DenseElementsAttr GetShape(Value output_val, bool truncate = false) {
  auto output_shape = output_val.getType().dyn_cast<ShapedType>().getShape();

  SmallVector<int32_t> shape;
  shape.reserve(output_shape.size());

  TF_EXPECT_OK(ShapeAnnotationsMatch(*graph, shape_info, expected));
}

// Test that shape inference uses user-given `arg_shapes` to propagate shapes.
TEST(ShapeInferenceTest, UseArgShapesForVariableBatchSize) {
  Scope root = Scope::NewRootScope().ExitOnError();
  auto a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT,
                            ops::Placeholder::Shape({-1, 3}));
  auto b = ops::Placeholder(root.WithOpName("B"), DT_FLOAT,

// Eg 1. An update in canonical form:
//  * indices shape(A,B,C)
//  * updates shape(A,B,D,E,F)
// Then:
//  * D,E,F are the update window dims [2,3,4]
//  * C is the index vector dimension
//  * A,B iterate over the updates and indices
//
// If `update_window_dims` are not the trailing dimensions then updates must be
// transposed.
//
// Eg 2. An update in non-canonical form:
//  * indices shape(a,b,c)
//  * updates shape(d,e,f,a,b)

  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();
}

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)]>;

# TODO(fengliuai): Figure out what to do here. A "golden" end-to-end test is overly fragile.

    if (BroadcastVector<double>(shaped.getDimSize(quant_dim), scales) ||
        BroadcastVector<int64_t>(shaped.getDimSize(quant_dim), zero_points)) {
      return {};
    }
  } else if ((new_shape.size() == shape.size() + 1) && new_shape.front() == 1) {
    // Handle the [A, B, C] -> [1, A, B, C] reshape case.
    if (!(std::equal(shape.begin(), shape.end(), new_shape.begin() + 1) &&