EXPECT_TRUE(IsStablehloOp(*constant_op));
}

TEST_F(StablehloTypeUtilsTest, IsStablehloOpFailsWithArithOp) {
  const OwningOpRef<mlir::arith::ConstantOp> constant_op =
      builder_.create<mlir::arith::ConstantOp>(builder_.getUnknownLoc(),
                                               builder_.getI32IntegerAttr(0));
  EXPECT_FALSE(IsStablehloOp(*constant_op));
}

}  // namespace

#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo

namespace mlir::quant::stablehlo {

// Checks if an op is from StableHLO dialect.
inline bool IsStablehloOp(Operation* op) {
  return op->getDialect()->getNamespace() ==
         mlir::stablehlo::StablehloDialect::getDialectNamespace();
}

}  // namespace mlir::quant::stablehlo