Build uploaded using `kernels`.

.gitattributes

@@ -35,3 +35,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 build/torch28-metal-aarch64-darwin/mlx_rmsnorm/_mlx_rmsnorm_97571a8_dirty.abi3.so filter=lfs diff=lfs merge=lfs -text
 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/__init__.py

@@ -1,25 +1,3 @@
-from ._ops import ops
+from .layers import RMSNorm, rmsnorm_forward, rmsnorm_backward
 
-import torch
-
-
-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
+__all__ = ["RMSNorm", "rmsnorm_forward", "rmsnorm_backward"]

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

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

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

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

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

@@ -0,0 +1,32 @@
+from ._ops import ops
+
+import torch
+
+class RMSNorm(torch.nn.Module):
+    weight: torch.Tensor
+    eps: float
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        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,25 +1,3 @@
-from ._ops import ops
+from .layers import RMSNorm, rmsnorm_forward, rmsnorm_backward
 
-import torch
-
-
-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
+__all__ = ["RMSNorm", "rmsnorm_forward", "rmsnorm_backward"]

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

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

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

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

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

@@ -0,0 +1,32 @@
+from ._ops import ops
+
+import torch
+
+class RMSNorm(torch.nn.Module):
+    weight: torch.Tensor
+    eps: float
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        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