Move `RMSNorm` to `layers`

torch-ext/triton_layer_norm/__init__.py

@@ -1,3 +1,5 @@
-from .layer_norm import RMSNorm, layer_norm_fn, layer_norm_linear_fn, rms_norm_fn
+from .layer_norm import layer_norm_fn, layer_norm_linear_fn, rms_norm_fn
 
-__all__ = ["RMSNorm", "layer_norm_fn", "layer_norm_linear_fn", "rms_norm_fn"]
+from . import layers
+
+__all__ = ["layers", "layer_norm_fn", "layer_norm_linear_fn", "rms_norm_fn"]

torch-ext/triton_layer_norm/layers.py

@@ -0,0 +1,4 @@
+from .layer_norm import RMSNorm
+
+
+__all__ = ["RMSNorm"]