Merges a tuple (s, V) into a single `torch.Tensor`, where the vector channels are flattened and appended to the scalar channels. Should be used only if the tuple representation cannot be used. Use `_split(x, nv)` to reverse.
(s, v)
| 77 | |
| 78 | |
| 79 | def _merge(s, v): |
| 80 | ''' |
| 81 | Merges a tuple (s, V) into a single `torch.Tensor`, where the |
| 82 | vector channels are flattened and appended to the scalar channels. |
| 83 | Should be used only if the tuple representation cannot be used. |
| 84 | Use `_split(x, nv)` to reverse. |
| 85 | ''' |
| 86 | v = torch.reshape(v, v.shape[:-2] + (3*v.shape[-2],)) |
| 87 | return torch.cat([s, v], -1) |
| 88 | |
| 89 | |
| 90 | class GVP(nn.Module): |