| 869 | def test_compile_fix_broken_ops(self) -> None: |
| 870 | class ExportableLoop(nn.Module): |
| 871 | def __init__(self, hidden_size, out_channels): |
| 872 | super().__init__() |
| 873 | self.hidden_size = hidden_size |
| 874 | self.B = nn.Parameter(torch.randn(hidden_size, 1)) # (H, in_channels) |
| 875 | self.C = nn.Parameter( |
| 876 | torch.randn(out_channels, hidden_size) |
| 877 | ) # (C_out, H) |
| 878 | A = torch.randn(2, hidden_size) |
| 879 | self.A_real = nn.Parameter(A[0].clone()) |
| 880 | self.A_imag = nn.Parameter(A[1].clone()) |
| 881 | |
| 882 | def update_state(self, h, x_t): |
| 883 | # h: [B, 2, H], x_t: [B, H] |