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;

      // 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);