Build

build/torch-universal/kernels_test/__init__.py

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

build/torch-universal/kernels_test/layers.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+
+
+class LinearImplicitBackward(nn.Module):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.linear(input, self.weight, self.bias)
+
+
+class LinearBackward(nn.Module):
+    has_backward = True
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.linear(input, self.weight, self.bias)
+
+
+class LinearNoBackward(nn.Module):
+    has_backward = False
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.linear(input, self.weight, self.bias)
+
+
+__all__ = ["LinearImplicitBackward", "LinearBackward", "LinearNoBackward"]