| 89 | |
| 90 | |
| 91 | class FTMapDynShape(Module): |
| 92 | def __init__(self): |
| 93 | super().__init__() |
| 94 | |
| 95 | def forward(self, xs, y): |
| 96 | def f(x, y): |
| 97 | return x + y |
| 98 | |
| 99 | return torch_map(f, xs, y) + xs |
| 100 | |
| 101 | def get_upper_bound_inputs(self): |
| 102 | return torch.rand(4, 4), torch.rand(4) |
| 103 | |
| 104 | def get_random_inputs(self): |
| 105 | return torch.rand(2, 4), torch.rand(4) |
no outgoing calls