| 69 | return relax.op.nn.relu(state) |
| 70 | |
| 71 | class Layer(nn.Module): |
| 72 | define_subroutine = True |
| 73 | |
| 74 | def __init__(self, in_features, out_features): |
| 75 | self.weights = nn.Parameter( |
| 76 | (in_features, out_features), dtype="float32", name="weights" |
| 77 | ) |
| 78 | self.activation = Activation() |
| 79 | |
| 80 | def forward(self, input: relax.Expr) -> relax.Var: |
| 81 | state = relax.op.matmul(input, self.weights) |
| 82 | return self.activation(state) |
| 83 | |
| 84 | @I.ir_module |
| 85 | class Expected: |
no outgoing calls
searching dependent graphs…