hyper-connections 0.4.5__py3-none-any.whl → 0.4.7__py3-none-any.whl

hyper_connections/mHCv2.py

@@ -47,6 +47,10 @@ def l1norm(t, dim):
     return F.normalize(t, p = 1, dim = dim)
 
 def sinkhorn_knopps(log_alpha, iters = 20):
+
+    if iters <= 0:
+        return log_alpha
+
     assert log_alpha.shape[-2] == log_alpha.shape[-1]
 
     dtype = log_alpha.dtype
@@ -63,6 +67,10 @@ def sinkhorn_knopps(log_alpha, iters = 20):
     return alpha.to(dtype)
 
 def log_domain_sinkhorn_knopps(log_alpha, iters = 20):
+
+    if iters <= 0:
+        return log_alpha
+
     assert log_alpha.shape[-2] == log_alpha.shape[-1]
 
     dtype = log_alpha.dtype
@@ -116,8 +124,8 @@ def get_init_and_expand_reduce_stream_functions(
 
     hyper_conn_klass = ManifoldConstrainedHyperConnections if not disable else Residual
 
-    kwargs.pop('add_attn_pool_reduce_stream', None)
     init_hyper_conn_fn = partial(hyper_conn_klass, num_streams, num_fracs = num_fracs, sinkhorn_iters = sinkhorn_iters, use_triton_sinkhorn = use_triton_sinkhorn, **kwargs)
+
     expand_reduce_fns = get_expand_reduce_stream_functions(
         num_streams,
         add_stream_embed = add_stream_embed,

hyper_connections/vit.py

@@ -5,7 +5,7 @@ from torch.nn import Module, ModuleList
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
-from hyper_connections.manifold_constrained_hyper_connections import mHC
+from hyper_connections.mHCv2 import mHC
 
 # helpers
 
@@ -66,12 +66,12 @@ class Attention(Module):
         return self.to_out(out)
 
 class Transformer(Module):
-    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., num_residual_streams = 4, num_dynamic_alpha_proposals = 1):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., num_residual_streams = 4, mhc_kwargs = dict()):
         super().__init__()
         self.norm = nn.LayerNorm(dim)
         self.layers = ModuleList([])
 
-        init_hyper_conn, self.expand_streams, self.reduce_streams = mHC.get_init_and_expand_reduce_stream_functions(num_residual_streams, num_dynamic_alpha_proposals = num_dynamic_alpha_proposals)
+        init_hyper_conn, self.expand_streams, self.reduce_streams = mHC.get_init_and_expand_reduce_stream_functions(num_residual_streams, **mhc_kwargs)
 
         for _ in range(depth):
             self.layers.append(ModuleList([
@@ -92,7 +92,7 @@ class Transformer(Module):
         return self.norm(x)
 
 class ViT(Module):
-    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., num_residual_streams = 4, num_dynamic_alpha_proposals = 1):
+    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., num_residual_streams = 4, mhc_kwargs = dict(num_dynamic_alpha_proposals = 1)):
         super().__init__()
         image_height, image_width = pair(image_size)
         patch_height, patch_width = pair(patch_size)
@@ -117,7 +117,7 @@ class ViT(Module):
 
         self.dropout = nn.Dropout(emb_dropout)
 
-        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, num_residual_streams, num_dynamic_alpha_proposals)
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout, num_residual_streams, mhc_kwargs)
 
         self.pool = pool
         self.to_latent = nn.Identity()
@@ -154,7 +154,10 @@ if __name__ == '__main__':
         mlp_dim = 2048,
         dropout = 0.1,
         emb_dropout = 0.1,
-        num_residual_streams = 4
+        num_residual_streams = 4,
+        mhc_kwargs = dict(
+            use_triton_sinkhorn = False
+        )
     )
 
     img = torch.randn(1, 3, 256, 256)

{hyper_connections-0.4.5.dist-info → hyper_connections-0.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.4.5
+Version: 0.4.7
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections

{hyper_connections-0.4.5.dist-info → hyper_connections-0.4.7.dist-info}/RECORD

@@ -3,12 +3,12 @@ hyper_connections/hyper_connections.py,sha256=2F-104cGE82KCK0KeC07NSOJNPT-0PCtvX
 hyper_connections/hyper_connections_channel_first.py,sha256=5vAen4WXxNI9K07ndLBQJwdJv-OjoXznta5EIQTaQNA,6512
 hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=yn2AlFB6qCYQeRhhhaMlCM3mxxLEdWCYwU2p9TsMwWI,7835
 hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=iFPw5pgCRHTo16nBJ2PExKSrvTyCh7ba7Py14P1oSPE,11311
-hyper_connections/mHCv2.py,sha256=wCtp87OFI3QfosdSL-1qwsiQN9f8gX32_0r8GQGO7P0,17411
+hyper_connections/mHCv2.py,sha256=XB2HwxTo7daZvy9fzF8UjHI12ephwgE91h9AH2Ou4WI,17452
 hyper_connections/manifold_constrained_hyper_connections.py,sha256=E4os-6q_SMjJO1JD0EG8rFTCXA7MQoy-aqUlM7KVS5Q,18269
 hyper_connections/residuals.py,sha256=JVSFJj_H7xQ3_Fd-pZH5Hdv9SveAQu29jQNvMyom5ek,921
 hyper_connections/triton_sinkhorn.py,sha256=n2WyQcUemtv5T5Sk2nljnSpV2hEED4I3HaPsIUy4638,5905
-hyper_connections/vit.py,sha256=BOWVfCAIzDQdnTq8OBzNUyiKGGILYZkIQ6mr1GKJVB0,5225
-hyper_connections-0.4.5.dist-info/METADATA,sha256=sWVb_-yVRmxL8AsAPsk0VdRXOa25uG9zKNc8S_oAXg8,6704
-hyper_connections-0.4.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hyper_connections-0.4.5.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.4.5.dist-info/RECORD,,
+hyper_connections/vit.py,sha256=dh8AVMUPaUHuWxXJEHoMW_G5nj-EQQjDmgbPwwhiq5g,5215
+hyper_connections-0.4.7.dist-info/METADATA,sha256=2ajn-IuCxuUjgnOw5dEBxXKqLJbyQohHSGgKJ2dZFoA,6704
+hyper_connections-0.4.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hyper_connections-0.4.7.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.4.7.dist-info/RECORD,,