update to mHC

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.14.5"
+version = "0.15.0"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -29,7 +29,7 @@ dependencies = [
     'einx>=0.3.0',
     'einops>=0.8.0',
     'ema-pytorch',
-    'hyper-connections>=0.0.10',
+    'hyper-connections>=0.3.6',
     'jaxtyping',
     'loguru',
     'rotary_embedding_torch>=0.8.4',

transfusion_pytorch/transfusion.py

@@ -44,7 +44,7 @@
 
 from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
 
-from hyper_connections import HyperConnections
+from hyper_connections import ManifoldConstrainedHyperConnections
 
 from tqdm import tqdm
 from loguru import logger
@@ -1067,7 +1067,7 @@ def __init__(
 
         counter = count()
 
-        init_residual_fn, self.expand_stream, self.reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(num_residual_streams, num_fracs = num_residual_fracs)
+        init_residual_fn, self.expand_stream, self.reduce_stream = ManifoldConstrainedHyperConnections.get_init_and_expand_reduce_stream_functions(num_residual_streams, num_fracs = num_residual_fracs)
 
         # layers