A simple version of `jt.repeat_interleave` for duplicating a matrix while interleaving the copy.
(m)
| 32 | return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)') |
| 33 | |
| 34 | def duplicate_interleave(m): |
| 35 | """ |
| 36 | A simple version of `jt.repeat_interleave` for duplicating a matrix while interleaving the copy. |
| 37 | """ |
| 38 | dim0 = m.shape[0] |
| 39 | m = m.view(-1, 1) # flatten the matrix |
| 40 | m = m.repeat(1, 2) # repeat all elements into the 2nd dimension |
| 41 | m = m.view(dim0, -1) # reshape into a matrix, interleaving the copy |
| 42 | return m |
| 43 | |
| 44 | |
| 45 | def apply_rotary_pos_emb(x, sincos, offset=0): |
no outgoing calls
no test coverage detected