LogicalResult rewriteOpWithDynamicInput(Operation* op,
                                          PatternRewriter& rewriter) const {
    auto cond = op->getOperand(0);
    auto lhs = op->getOperand(1);
    auto rhs = op->getOperand(2);
    auto out = op->getResult(0);

    // Matches fail when lhs|rhs|cond is unranked tensor.
    // TODO(b/176202543): Support unranked tensor.

    if (!erf_user_user_input) return failure();

    // Check `mul 0.5` and `mul 1/sqrt(2)` refer to the same input.
    if (erf_user_user_input->getOperand(0) != erf_input->getOperand(0)) {
      return failure();
    }

    // Check the structural matches have the correct op type and values.
    auto rhs_mul = llvm::dyn_cast_or_null<stablehlo::MulOp>(erf_input);

      Or<[TFL_OperandIsUnrankedPred<0>,
          CPred<"$_op.getOperand(0).getType().cast<ShapedType>().getRank() <= 1">,
          CPred<"$_op.getOperand(0).getType().cast<ShapedType>().getRank() == 2 && !$_op.getOperand(0).getType().cast<ShapedType>().hasStaticShape()">,
          CPred<"$_op.getOperand(0).getType().cast<ShapedType>().getRank() == 2 && $_op.getOperand(0).getType().cast<ShapedType>().getShape()[1] <= 4">]>>]> {

  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];
  output_shape[3] = composite_result_shape[1];

  auto input_type = mlir::cast<ShapedType>(old_op->getOperand(0).getType());

  return RankedTensorType::get(output_shape, input_type.getElementType());
}
}  // namespace odml

    }
  }

  if (!is_supported_affine_op) return failure();

  // Bind input, output and weight to the given values.
  binding_output = op->getResult(0);
  binding_input = op->getOperand(0);
  binding_weight = op->getOperand(1);
  return success();
}

// Makes the 1D value broadcastable with the `rhs_shape`.
Value MakeOneDimValueBroadcastable(OpBuilder& builder, Location loc,

bool GetValueAsConstant(Value val, AttrT &attr) {
  while (auto result = mlir::dyn_cast<OpResult>(val)) {
    Operation *op = result.getOwner();
    if (!isa<IdentityOp>(op) && !isa<IdentityNOp>(op)) break;
    val = op->getOperand(result.getResultNumber());
  }
  return matchPattern(val, m_Constant(&attr));
}

// Checks if both compilation and replication attributes are present in the
// operation, and if their values are valid.

                                            TypeConverter &converter,
                                            IRRewriter &rewriter) {
  for (int i = 0; i < op->getNumOperands(); ++i) {
    auto operand = op->getOperand(i);
    auto new_operand = converter.materializeArgumentConversion(
        rewriter, op->getLoc(), converter.convertType(operand.getType()),
        {operand});
    op->setOperand(i, new_operand);
  }
}

TorchAvgPoolData GetTorchAvgPoolData(CompositeOp op) {
  auto composite_attrs = op.getCompositeAttributes();
  TorchAvgPoolData data;

  auto op_type = mlir::cast<RankedTensorType>(op.getOperand(0).getType());

  data.n = op_type.getShape()[0];
  data.c = op_type.getShape()[1];
  data.h_in = op_type.getShape()[2];
  data.w_in = op_type.getShape()[3];

  std::vector<int32_t> kernel_size;

      mlir::Value return_val = return_op->getOperand(idx);
      if (auto owner = mlir::dyn_cast_or_null<mlir::tensor::CastOp>(
              return_val.getDefiningOp())) {
        // For bounded dynamic type, get a static size by taking bounds as the
        // dimensions. These dimensions are marked as dynamic in xla::Shape
        // below.
        buffer_ty = GetBufferType(owner.getOperand().getType());

  auto op_name = std::move(op_name_or_status).value();
  if (op_name == "NextIteration")
    return mlir::isa<mlir::TF::TensorFlowRefType>(
        mlir::getElementTypeOrSelf(op->getOperand(0).getType()));

  if (op_name == "Enter" || op_name == "Exit" || op_name == "Switch" ||
      op_name == "Merge") {
    return mlir::isa<mlir::TF::TensorFlowRefType>(