Build uploaded using `kernels`.

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

@@ -7,16 +7,19 @@ def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5
     original_shape = x.shape
     x = x.view(-1, x.shape[-1])
     weight = weight.view(-1)
-    out = ops.launch_forward_kernel(x, weight, epsilon)
-    out = out.view(original_shape)
-    return out
+    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, grad_weight = ops.launch_backward_kernel(x, weight, grad_output, epsilon)
+    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/_mlx_rmsnorm_97571a8_dirty.abi3.so

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

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

@@ -7,16 +7,19 @@ def rmsnorm_forward(x: torch.Tensor, weight: torch.Tensor, epsilon: float = 1e-5
     original_shape = x.shape
     x = x.view(-1, x.shape[-1])
     weight = weight.view(-1)
-    out = ops.launch_forward_kernel(x, weight, epsilon)
-    out = out.view(original_shape)
-    return out
+    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, grad_weight = ops.launch_backward_kernel(x, weight, grad_output, epsilon)
+    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/_mlx_rmsnorm_97571a8_dirty.abi3.so

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