auto layer_stats =
      mlir::cast<mlir::DenseFPElementsAttr>(stats_op.getLayerStats());
  std::optional<mlir::ElementsAttr> axis_stats = stats_op.getAxisStats();
  std::optional<uint64_t> axis = stats_op.getAxis();
  std::vector<float> mins, maxs;
  mlir::DenseFPElementsAttr min_max_attr =
      axis_stats.has_value()
          ? mlir::cast<mlir::DenseFPElementsAttr>(axis_stats.value())
          : layer_stats;

//===----------------------------------------------------------------------===//

OpFoldResult PowOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  auto constant_y = operands[1].dyn_cast_or_null<DenseFPElementsAttr>();
  if (constant_y && constant_y.isSplat()) {
    APFloat y_value = constant_y.getSplatValue<APFloat>();
    auto output_type = getType().cast<ShapedType>();
    if (y_value.isZero() && output_type.hasStaticShape()) {

      // scratch1 = rsqrt(var + epsilon)
      RankedTensorType scalar_float =
          tensorflow::GetTypeFromTFTensorShape({}, kernel_type);
      auto epsilon = rewriter.create<ConstantOp>(
          loc, DenseFPElementsAttr::get(scalar_float, {op.getEpsilon()}));
      auto add_op = rewriter.create<chlo::BroadcastAddOp>(
          loc, var, epsilon.getResult(), scalar_broadcast_dims);