chore(poly-norm): add ROCm build artifacts

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text

build/torch26-cxx11-rocm62-x86_64-linux/activation/__init__.py

@@ -0,0 +1,21 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/torch26-cxx11-rocm62-x86_64-linux/activation/_activation_44e9845_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:527e5aac540e24dc3791fd423fea23f687ea3cffdb627c6a6e35f4df1aa7dec4
+size 2460736

build/torch26-cxx11-rocm62-x86_64-linux/activation/_ops.py

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

build/torch26-cxx11-rocm62-x86_64-linux/activation/layers.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+
+from .poly_norm import PolyNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)

build/torch26-cxx11-rocm62-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None

build/torch27-cxx11-rocm63-x86_64-linux/activation/__init__.py

@@ -0,0 +1,21 @@
+import torch
+
+from . import layers
+from ._ops import ops
+from .poly_norm import PolyNormFunction
+
+
+def poly_norm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    return PolyNormFunction.apply(x, weight, bias, eps)
+
+
+__all__ = [
+    "poly_norm",
+    "layers",
+    "ops",
+]

build/torch27-cxx11-rocm63-x86_64-linux/activation/_activation_44e9845_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3e00863b72834e1d121e377e41724b1479703051aed7d9d8a64019d6a92bf54
+size 2447432

build/torch27-cxx11-rocm63-x86_64-linux/activation/_ops.py

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

build/torch27-cxx11-rocm63-x86_64-linux/activation/layers.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+
+from .poly_norm import PolyNormFunction
+
+
+class PolyNorm(nn.Module):
+    def __init__(self, eps):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(3) / 3)
+        self.bias = torch.nn.Parameter(torch.zeros(1))
+        self.eps = eps
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ):
+        return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)

build/torch27-cxx11-rocm63-x86_64-linux/activation/poly_norm.py

@@ -0,0 +1,41 @@
+import torch
+
+from ._ops import ops
+
+
+# Inherit from Function
+class PolyNormFunction(torch.autograd.Function):
+    # Note that forward, setup_context, and backward are @staticmethods
+    @staticmethod
+    def forward(input, weight, bias, eps):
+        output = torch.empty_like(input)
+        ops.poly_norm(output, input, weight, bias, eps)
+        return output
+
+    @staticmethod
+    # inputs is a Tuple of all of the inputs passed to forward.
+    # output is the output of the forward().
+    def setup_context(ctx, inputs, output):
+        input, weight, bias, eps = inputs
+        ctx.save_for_backward(input, weight)
+        ctx.eps = eps
+
+    # This function has only a single output, so it gets only one gradient
+    @staticmethod
+    def backward(ctx, output_grad):
+        input, weight = ctx.saved_tensors
+        eps = ctx.eps
+
+        input_grad = torch.empty_like(input) if ctx.needs_input_grad[0] else None
+        weight_grad = torch.empty_like(weight) if ctx.needs_input_grad[1] else None
+        bias_grad = (
+            torch.empty(1, dtype=weight.dtype, device=weight.device)
+            if ctx.needs_input_grad[2]
+            else None
+        )
+
+        ops.poly_norm_backward(
+            input_grad, weight_grad, bias_grad, output_grad, input, weight, eps
+        )
+
+        return input_grad, weight_grad, bias_grad, None