Add layer

README.md

@@ -0,0 +1,3 @@
+---
+license: apache-2.0
+---

build.toml

@@ -0,0 +1,3 @@
+[general]
+name = "state_test"
+universal = true

build/torch-universal/state_test/__init__.py

@@ -0,0 +1,3 @@
+from . import layers
+
+__all__ = ["layers"]

build/torch-universal/state_test/_ops.py

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

build/torch-universal/state_test/layers.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+@dataclass
+class ReLUState:
+    zeros: torch.Tensor
+
+
+class StatefulReLU(nn.Module):
+    can_torch_compile = True
+    has_backward = True
+
+    hidden_size: int
+
+    @staticmethod
+    def create_state(device: torch.device, layer: nn.Module) -> ReLUState:
+        zeros = torch.zeros(layer.hidden_size, device=device)
+        return ReLUState(zeros=zeros)
+
+    def forward_with_state(self, state: ReLUState, input: torch.Tensor) -> torch.Tensor:
+        return torch.maximum(input, state.zeros)
+
+
+__all__ = ["StatefulReLU"]

flake.lock

@@ -0,0 +1,168 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-compat_2": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "flake-utils_2": {
+      "inputs": {
+        "systems": "systems_2"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "hf-nix": {
+      "inputs": {
+        "flake-compat": "flake-compat_2",
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1757675377,
+        "narHash": "sha256-JQKZOI1ZYO4faJnanuoTXziSmqzXe5rEFSGliWDWqWw=",
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "rev": "faf3354403a7381958d08e826c15fe30f6986a4f",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "hf-nix": "hf-nix",
+        "nixpkgs": [
+          "kernel-builder",
+          "hf-nix",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1758103102,
+        "narHash": "sha256-z9E9FxuxuxUztG5DbUcOvKBHvd27gBY9617t9x2QE6M=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "94369928dc09ea7753c58495e3e406ac26f6c378",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1755963616,
+        "narHash": "sha256-6yD0ww/S8n+U2uPYcJZ3DRURP8Kx036GRpR2uPNZroE=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "73e96df7cff5783f45e21342a75a1540c4eddce4",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    },
+    "systems_2": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for kernels tests";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs {
+      path = ./.;
+      rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
+    };
+}

torch-ext/state_test/__init__.py

@@ -0,0 +1,3 @@
+from . import layers
+
+__all__ = ["layers"]

torch-ext/state_test/layers.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+@dataclass
+class ReLUState:
+    zeros: torch.Tensor
+
+
+class StatefulReLU(nn.Module):
+    can_torch_compile = True
+    has_backward = True
+
+    hidden_size: int
+
+    @staticmethod
+    def create_state(device: torch.device, layer: nn.Module) -> ReLUState:
+        zeros = torch.zeros(layer.hidden_size, device=device)
+        return ReLUState(zeros=zeros)
+
+    def forward_with_state(self, state: ReLUState, input: torch.Tensor) -> torch.Tensor:
+        return torch.maximum(input, state.zeros)
+
+
+__all__ = ["StatefulReLU"]