if (!input_ty.hasRank() || input_ty.getRank() != 4) {
      return failure();
    }

    int64_t batch_cst = input_ty.getShape()[0];
    int64_t channels_cst = input_ty.getShape()[3];

    int64_t in_y_cst = input_ty.getShape()[1];
    int64_t in_x_cst = input_ty.getShape()[2];
    int64_t in_spatial_cst =
        in_y_cst < 0 || in_x_cst < 0 ? -1 : in_y_cst * in_x_cst;

                           Location loc) {
  auto buffer_type = mlir::cast<RankedTensorType>(buffer.getType());
  if (buffer_type.getShape().size() == 1) return index;
  // Create a concat of index and trailing zeros.
  llvm::SmallVector<int64_t, 8> zeros(buffer_type.getShape().size() - 1, 0);
  auto zeros_tensor = GetR1Const(zeros, builder, loc);
  return builder.create<TF::ConcatV2Op>(
      loc,

  if (auto tensor_type = dyn_cast<RankedTensorType>(input_type))
    return RankedTensorType::get(tensor_type.getShape(), elemental_type);
  if (auto tensor_type = dyn_cast<UnrankedTensorType>(input_type))
    return UnrankedTensorType::get(elemental_type);
  if (auto vector_type = dyn_cast<VectorType>(input_type))
    return VectorType::get(vector_type.getShape(), elemental_type);

  // If the expressed types match, just use the new elemental type.

      SmallVector<int64_t> new_shape(shape.begin(), shape.end());
      // If dimension of the input type is dynamic. Update the
      // bounds of the dim with the new type if needed.
      for (int i = 0; i < input_ty.getShape().size(); i++) {
        if (hlo::isDynamicDimSize(input_ty.getShape()[i])) {
          new_bounds[i] = new_shape[i];
          new_shape[i] = ShapedType::kDynamic;
        }
      }

  // Extract Z from %dot.
  ArrayRef<int64_t> shape_z =
      dot.getType().getShape().drop_front(shape_b.size() + shape_y2.size());

  // Check %after shape.
  if (reshape_after.getType().getShape() !=
      ArrayRef<int64_t>(llvm::to_vector(
          llvm::concat<const int64_t>(shape_b, shape_y1, shape_z)))) {
    return failure();
  }

  rewriter.replaceOpWithNewOp<mhlo::DotGeneralOp>(

  RankedTensorType tensorType = mlir::cast<RankedTensorType>(value.getType());
  Type element_type = tensorType.getElementType();

  int rank = tensorType.getShape().size();
  int num_rows = tensorType.getShape()[rank - 2];
  int num_cols = tensorType.getShape()[rank - 1];

  std::vector<Value> sliced;

  if (batch_size == 1) {
    // Batch size is 1, no splitting is required

      if (input_ty_i.hasStaticShape() &&
          input_ty_i.getShape() != input_ty_0.getShape()) {
        return op.emitOpError()
               << "inputs[" << i << "] has shape [" << input_ty_i.getShape()
               << "] different than the shape of inputs[0]: "
               << input_ty_0.getShape();
      }
    }

    if (op.getDimensionsToReduce().size() > input_ty_0.getRank()) {

  if (!type.hasRank()) return type;

  auto shape = type.getShape();
  SmallVector<int64_t, 4> new_shape(shape.size());

  for (int i = 0; i < permutation.size(); ++i) {
    int64_t index = permutation[i];
    assert(index < shape.size());
    new_shape[index] = shape[i];
  }

  return type.clone(new_shape);
}

// Move Transpose operations that permute `op` operands after the `op`.

                                   ConversionPatternRewriter& rewriter) {
  if (index_vector_dim == indices_type.getRank()) {
    llvm::SmallVector<int64_t, 4> new_start_indices_shape(
        indices_type.getShape().begin(), indices_type.getShape().end());
    new_start_indices_shape.push_back(1);
    indices_type = RankedTensorType::get(new_start_indices_shape,
                                         indices_type.getElementType());

def GetSqueezedPermutation: NativeCodeCall<"GetSqueezedPermutation($0, $1)">;

// Check to see if the tensor dimensions can be Squeezed by eliminating 1s'
def CanSqueezeTensor : Constraint<CPred<
  "GetShape($0).getNumElements() > GetSqueezedShape($0).getNumElements()">>;


// Pattern to convert TFL_TransposeOp with rank>6 to rank<=6 if there are
// redundant dimensions in the tensor. For example- [2x1x3] == [2x3] and 1 is