const bool is_per_axis_quantizable =
      IsDotGeneralFullyConnected(dot_general_op).value();
  if (!is_per_axis_quantizable) return std::nullopt;
  return filter_rank - 1;
}

bool ContainsConvOrDot(StringRef str) {
  return str.contains("_conv") || str.contains("_dot_general");
}

  ASSERT_THAT(main_fn, NotNull());

  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
  const StringRef function_name = GetEntryFunctionName(call_op);
  EXPECT_TRUE(ContainsConvOrDot(function_name));
}

TEST_F(AttrsAndConstraintsTest, ContainsConvOrDotFalse) {
  OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kModuleXlaCallModuleNoEntryNoQuantTrait);
  ASSERT_TRUE(module_op);

std::optional<int64_t> GetDotGeneralQuantizationDim(
    ::mlir::stablehlo::DotGeneralOp dot_general_op);

// Checks if a `StringRef` contains 'conv' or 'dot_general'.
bool ContainsConvOrDot(StringRef str);

}  // namespace mlir::quant

  if (auto call_op = dyn_cast<TF::XlaCallModuleOp>(op)) {
    StringRef entry_function_name = GetEntryFunctionName(call_op);
    absl::StatusOr<Method> quantization_method = GetQuantizationMethod(call_op);
    return ContainsConvOrDot(entry_function_name) && quantization_method.ok() &&
           quantization_method->has_weight_only_ptq();
  }
  return false;
}

SmallVector<func::FuncOp> GetSortedFunctions(ModuleOp module_op) {