bn_op.getLoc(), mean_var_type, square_diff, reduce_dim_op,
        /*keep_dims=*/rewriter.getBoolAttr(false));

    // Invoke BatchNormInferenceOp
    ::mlir::FloatAttr epsilon = bn_op.getEpsilonAttr();
    ::mlir::Value batch_norm = rewriter.create<mhlo::BatchNormInferenceOp>(
        bn_op.getLoc(), inputs, bn_op.getScale(), bn_op.getOffset(), mean,
        variance, epsilon, rewriter.getI64IntegerAttr(feature_index));