Build uploaded using `kernels`.

.gitattributes

@@ -37,3 +37,5 @@ build/torch28-metal-aarch64-darwin/mlx_rmsnorm/_mlx_rmsnorm_97571a8_dirty.abi3.s
 build/torch29-metal-aarch64-darwin/mlx_rmsnorm/_mlx_rmsnorm_97571a8_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch28-metal-aarch64-darwin/mlx_rmsnorm/_mlx_rmsnorm_57d865f.abi3.so filter=lfs diff=lfs merge=lfs -text
 build/torch29-metal-aarch64-darwin/mlx_rmsnorm/_mlx_rmsnorm_57d865f.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch28-metal-aarch64-darwin/mlx_rmsnorm/_mlx_rmsnorm_06e17fa.abi3.so filter=lfs diff=lfs merge=lfs -text
+build/torch29-metal-aarch64-darwin/mlx_rmsnorm/_mlx_rmsnorm_06e17fa.abi3.so filter=lfs diff=lfs merge=lfs -text

build/torch28-metal-aarch64-darwin/mlx_rmsnorm/__init__.py

@@ -1,3 +1,3 @@
-from .layers import RMSNorm, rmsnorm_forward, rmsnorm_backward
-
-__all__ = ["RMSNorm", "rmsnorm_forward", "rmsnorm_backward"]
+from .functions import rmsnorm_forward, rmsnorm_backward
+from . import layers
+__all__ = ["layers", "rmsnorm_forward", "rmsnorm_backward"]

build/torch28-metal-aarch64-darwin/mlx_rmsnorm/{_mlx_rmsnorm_57d865f.abi3.so → _mlx_rmsnorm_06e17fa.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:22dcc9b26a28ff2c1cc34697bd25bf55eeb9cd9d4f1eb32a352ed2343fecbe6d
+oid sha256:6064021fd19860d705283dfb31cd15e6e4e1776dc4c1af0e29c8ee421deafa64
 size 219168

build/torch28-metal-aarch64-darwin/mlx_rmsnorm/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _mlx_rmsnorm_57d865f
-ops = torch.ops._mlx_rmsnorm_57d865f
+from . import _mlx_rmsnorm_06e17fa
+ops = torch.ops._mlx_rmsnorm_06e17fa
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mlx_rmsnorm_57d865f::{op_name}"
+    return f"_mlx_rmsnorm_06e17fa::{op_name}"

build/torch28-metal-aarch64-darwin/mlx_rmsnorm/functions.py

@@ -0,0 +1,23 @@
+import torch
+from ._ops import ops
+
+def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    output = torch.zeros_like(x)
+    ops.launch_forward_kernel(x, weight, output, epsilon)
+    output = output.view(original_shape)
+    return output
+
+def rmsnorm_backward(x: torch.Tensor, weight: torch.Tensor, grad_output: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    grad_output = grad_output.view(-1)
+    grad_input = torch.zeros_like(x)
+    grad_weight = torch.zeros_like(weight)
+    ops.launch_backward_kernel(x, weight, grad_output, grad_input, grad_weight, epsilon)
+    grad_input = grad_input.view(original_shape)
+    grad_weight = grad_weight.view(original_shape)
+    return grad_input, grad_weight

build/torch28-metal-aarch64-darwin/mlx_rmsnorm/layers.py

@@ -1,5 +1,4 @@
-from ._ops import ops
-
+from .functions import rmsnorm_forward
 import torch
 
 class RMSNorm(torch.nn.Module):
@@ -10,23 +9,3 @@ class RMSNorm(torch.nn.Module):
         return rmsnorm_forward(x, self.weight, self.eps)
 
 
-def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
-    original_shape = x.shape
-    x = x.view(-1, x.shape[-1])
-    weight = weight.view(-1)
-    output = torch.zeros_like(x)
-    ops.launch_forward_kernel(x, weight, output, epsilon)
-    output = output.view(original_shape)
-    return output
-
-def rmsnorm_backward(x: torch.Tensor, weight: torch.Tensor, grad_output: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
-    original_shape = x.shape
-    x = x.view(-1, x.shape[-1])
-    weight = weight.view(-1)
-    grad_output = grad_output.view(-1)
-    grad_input = torch.zeros_like(x)
-    grad_weight = torch.zeros_like(weight)
-    ops.launch_backward_kernel(x, weight, grad_output, grad_input, grad_weight, epsilon)
-    grad_input = grad_input.view(original_shape)
-    grad_weight = grad_weight.view(original_shape)
-    return grad_input, grad_weight

build/torch29-metal-aarch64-darwin/mlx_rmsnorm/__init__.py

@@ -1,3 +1,3 @@
-from .layers import RMSNorm, rmsnorm_forward, rmsnorm_backward
-
-__all__ = ["RMSNorm", "rmsnorm_forward", "rmsnorm_backward"]
+from .functions import rmsnorm_forward, rmsnorm_backward
+from . import layers
+__all__ = ["layers", "rmsnorm_forward", "rmsnorm_backward"]

build/torch29-metal-aarch64-darwin/mlx_rmsnorm/{_mlx_rmsnorm_57d865f.abi3.so → _mlx_rmsnorm_06e17fa.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a507684a2ed7b127800e22962811212ac657ac126665f645e4a4bfe78e4bb3d8
+oid sha256:3eb0d24dc5c019fa77403103e936cff377e6003f2537fdfd0892f0a0f73ae922
 size 220128

build/torch29-metal-aarch64-darwin/mlx_rmsnorm/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _mlx_rmsnorm_57d865f
-ops = torch.ops._mlx_rmsnorm_57d865f
+from . import _mlx_rmsnorm_06e17fa
+ops = torch.ops._mlx_rmsnorm_06e17fa
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_mlx_rmsnorm_57d865f::{op_name}"
+    return f"_mlx_rmsnorm_06e17fa::{op_name}"

build/torch29-metal-aarch64-darwin/mlx_rmsnorm/functions.py

@@ -0,0 +1,23 @@
+import torch
+from ._ops import ops
+
+def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    output = torch.zeros_like(x)
+    ops.launch_forward_kernel(x, weight, output, epsilon)
+    output = output.view(original_shape)
+    return output
+
+def rmsnorm_backward(x: torch.Tensor, weight: torch.Tensor, grad_output: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    original_shape = x.shape
+    x = x.view(-1, x.shape[-1])
+    weight = weight.view(-1)
+    grad_output = grad_output.view(-1)
+    grad_input = torch.zeros_like(x)
+    grad_weight = torch.zeros_like(weight)
+    ops.launch_backward_kernel(x, weight, grad_output, grad_input, grad_weight, epsilon)
+    grad_input = grad_input.view(original_shape)
+    grad_weight = grad_weight.view(original_shape)
+    return grad_input, grad_weight

build/torch29-metal-aarch64-darwin/mlx_rmsnorm/layers.py

@@ -1,5 +1,4 @@
-from ._ops import ops
-
+from .functions import rmsnorm_forward
 import torch
 
 class RMSNorm(torch.nn.Module):
@@ -10,23 +9,3 @@ class RMSNorm(torch.nn.Module):
         return rmsnorm_forward(x, self.weight, self.eps)
 
 
-def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
-    original_shape = x.shape
-    x = x.view(-1, x.shape[-1])
-    weight = weight.view(-1)
-    output = torch.zeros_like(x)
-    ops.launch_forward_kernel(x, weight, output, epsilon)
-    output = output.view(original_shape)
-    return output
-
-def rmsnorm_backward(x: torch.Tensor, weight: torch.Tensor, grad_output: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
-    original_shape = x.shape
-    x = x.view(-1, x.shape[-1])
-    weight = weight.view(-1)
-    grad_output = grad_output.view(-1)
-    grad_input = torch.zeros_like(x)
-    grad_weight = torch.zeros_like(weight)
-    ops.launch_backward_kernel(x, weight, grad_output, grad_input, grad_weight, epsilon)
-    grad_input = grad_input.view(original_shape)
-    grad_weight = grad_weight.view(original_shape)
-    return grad_input, grad_weight