fix: sort feature_dropout correctly in __init__.py and simplify benchmark

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: factnn

benchmark/test_feature_dropout_perf.py

@@ -2,34 +2,17 @@
 import torch
 
 from benchmark.attri_util import FLOAT_DTYPES
-from benchmark.performance_utils import GenericBenchmarkExcluse1D
+from benchmark.performance_utils import GenericBenchmark2DOnly
 
 
 def feature_dropout_input_fn(shape, dtype, device):
     inp = torch.randn(shape, dtype=dtype, device=device)
     yield inp, 0.5, True
 
 
-class FeatureDropoutBenchmark(GenericBenchmarkExcluse1D):
-    def set_shapes(self, shape_file_path=None):
-        self.shapes = [
-            (8, 64),
-            (16, 128),
-            (32, 256),
-            (4, 64, 32, 32),
-            (8, 128, 16, 16),
-            (16, 256, 8, 8),
-            (2, 512, 14, 14),
-            (4, 1024, 7, 7),
-        ]
-
-    def set_more_shapes(self):
-        return None
-
-
 @pytest.mark.feature_dropout
 def test_perf_feature_dropout():
-    bench = FeatureDropoutBenchmark(
+    bench = GenericBenchmark2DOnly(
         input_fn=feature_dropout_input_fn,
         op_name="feature_dropout",
         torch_op=torch.feature_dropout,

src/flag_gems/__init__.py

@@ -199,10 +199,10 @@ def torch_ge(v):
     ("expm1_", expm1_),
     ("expm1.out", expm1_out),
     ("exponential_", exponential_),
-    ("feature_dropout", feature_dropout),
-    ("feature_dropout_", feature_dropout_),
     ("eye", eye),
     ("eye.m", eye_m),
+    ("feature_dropout", feature_dropout),
+    ("feature_dropout_", feature_dropout_),
     ("fill.Scalar", fill_scalar),
     ("fill.Scalar_out", fill_scalar_out),
     ("fill.Tensor", fill_tensor),