Splits a merged representation of (s, V) back into a tuple. Should be used only with `_merge(s, V)` and only if the tuple representation cannot be used. :param x: the `torch.Tensor` returned from `_merge` :param nv: the number of vector channels in the input to `_merge`
(x, nv)
| 63 | |
| 64 | |
| 65 | def _split(x, nv): |
| 66 | ''' |
| 67 | Splits a merged representation of (s, V) back into a tuple. |
| 68 | Should be used only with `_merge(s, V)` and only if the tuple |
| 69 | representation cannot be used. |
| 70 | |
| 71 | :param x: the `torch.Tensor` returned from `_merge` |
| 72 | :param nv: the number of vector channels in the input to `_merge` |
| 73 | ''' |
| 74 | v = torch.reshape(x[..., -3*nv:], x.shape[:-1] + (nv, 3)) |
| 75 | s = x[..., :-3*nv] |
| 76 | return s, v |
| 77 | |
| 78 | |
| 79 | def _merge(s, v): |