titans-pytorch 0.4.8__py3-none-any.whl → 0.4.9__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (6) hide show
  1. titans_pytorch/neural_memory.py +9 -0
  2. {titans_pytorch-0.4.8.dist-info → titans_pytorch-0.4.9.dist-info}/METADATA +1 -1
  3. titans_pytorch-0.4.9.dist-info/RECORD +8 -0
  4. titans_pytorch-0.4.8.dist-info/RECORD +0 -8
  5. {titans_pytorch-0.4.8.dist-info → titans_pytorch-0.4.9.dist-info}/WHEEL +0 -0
  6. {titans_pytorch-0.4.8.dist-info → titans_pytorch-0.4.9.dist-info}/licenses/LICENSE +0 -0
titans_pytorch/neural_memory.py CHANGED
@@ -164,6 +164,12 @@ def newtonschulz5(
164
164
  if t.ndim <= 3:
165
165
  return t
166
166
 
167
+ shape = t.shape
168
+ should_transpose = shape[2] > shape[-1]
169
+
170
+ if should_transpose:
171
+ t = t.transpose(-1, -2)
172
+
167
173
  t, inv_pack = pack_one_with_inverse(t, '* i j')
168
174
  t = t / t.norm(dim = (-1, -2), keepdim = True).clamp(min = eps)
169
175
 
@@ -174,6 +180,9 @@ def newtonschulz5(
174
180
  B = b * A + c * A @ A
175
181
  t = a * t + B @ t
176
182
 
183
+ if should_transpose:
184
+ t = t.transpose(-1, -2)
185
+
177
186
  return inv_pack(t)
178
187
 
179
188
  # multi head rmsnorm
{titans_pytorch-0.4.8.dist-info → titans_pytorch-0.4.9.dist-info}/METADATA RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: titans-pytorch
3
- Version: 0.4.8
3
+ Version: 0.4.9
4
4
  Summary: Titans
5
5
  Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
6
6
  Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
titans_pytorch-0.4.9.dist-info/RECORD ADDED
@@ -0,0 +1,8 @@
1
+ titans_pytorch/__init__.py,sha256=sVTOuRUkaIYabFExdLY6s1qXm1UwHHz_J19H8ZV-X74,338
2
+ titans_pytorch/mac_transformer.py,sha256=tz72141G5t3AOnxSVsOLtLptGtl8T7zROUvaTw2_XCY,26960
3
+ titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
4
+ titans_pytorch/neural_memory.py,sha256=JCK9t0dAYB6estqw9rrWENkI6qpsKF9QQf_MwXnWuJ0,34458
5
+ titans_pytorch-0.4.9.dist-info/METADATA,sha256=ony2yYgXUfdwP6QyM9o3BFSbdmH0HvSwA-BuClBacpQ,7873
6
+ titans_pytorch-0.4.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
7
+ titans_pytorch-0.4.9.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
8
+ titans_pytorch-0.4.9.dist-info/RECORD,,
titans_pytorch-0.4.8.dist-info/RECORD DELETED
@@ -1,8 +0,0 @@
1
- titans_pytorch/__init__.py,sha256=sVTOuRUkaIYabFExdLY6s1qXm1UwHHz_J19H8ZV-X74,338
2
- titans_pytorch/mac_transformer.py,sha256=tz72141G5t3AOnxSVsOLtLptGtl8T7zROUvaTw2_XCY,26960
3
- titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
4
- titans_pytorch/neural_memory.py,sha256=ypWXN8koY8pXt7IvlcMR1QM7cYJnWK_iYLEHy2pjx88,34277
5
- titans_pytorch-0.4.8.dist-info/METADATA,sha256=BbhF0oiPGdcgxrBGJziZvbvXUmS5lVlGVxpUGPmP0O8,7873
6
- titans_pytorch-0.4.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
7
- titans_pytorch-0.4.8.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
8
- titans_pytorch-0.4.8.dist-info/RECORD,,