Build

build/torch-universal/op_without_fake_test/__init__.py

@@ -0,0 +1,13 @@
+import torch
+import torch.nn.functional as F
+
+from . import layers
+
+
+@torch.library.custom_op("op_without_fake_test::silu_and_mul", mutates_args=())
+def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    d = x.shape[-1] // 2
+    return F.silu(x[..., :d]) * x[..., d:]
+
+
+["layers"]

build/torch-universal/op_without_fake_test/layers.py

@@ -0,0 +1,10 @@
+import torch
+import torch.nn as nn
+
+
+class SiluAndMul(nn.Module):
+    has_backward = False
+    can_torch_compile = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.op_without_fake_test.silu_and_mul(x)