hippoformer 0.0.9__tar.gz → 0.0.10__tar.gz

{hippoformer-0.0.9 → hippoformer-0.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hippoformer
-Version: 0.0.9
+Version: 0.0.10
 Summary: hippoformer
 Project-URL: Homepage, https://pypi.org/project/hippoformer/
 Project-URL: Repository, https://github.com/lucidrains/hippoformer

{hippoformer-0.0.9 → hippoformer-0.0.10}/hippoformer/hippoformer.py

@@ -455,6 +455,10 @@ class mmTEM(Module):
 
             update = self.assoc_scan(grad, expanded_beta.sigmoid(), momentum)
 
+            # store next momentum
+
+            next_momentum[key] = update[:, -1]
+
             # maybe muon
 
             if self.muon_update:
@@ -464,14 +468,13 @@ class mmTEM(Module):
 
             expanded_forget = repeat(forget, 'b t -> b t w', w = grad.shape[-1])
 
-            acc_update = self.assoc_scan(update, expanded_forget.sigmoid())
+            acc_update = self.assoc_scan(-update, expanded_forget.sigmoid(), param)
 
             acc_update = inverse_pack(acc_update)
 
             # set the next params and momentum, which can be passed back in
 
-            next_params[key] = param - acc_update[:, -1]
-            next_momentum[key] = update[:, -1]
+            next_params[key] =  acc_update[:, -1]
 
         # losses
 

{hippoformer-0.0.9 → hippoformer-0.0.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hippoformer"
-version = "0.0.9"
+version = "0.0.10"
 description = "hippoformer"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }