hyper-connections 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

hyper_connections/hyper_connections.py

@@ -11,6 +11,7 @@ import torch.nn.functional as F
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
+from einops.layers.torch import Reduce
 
 """
 ein notation:
@@ -35,11 +36,11 @@ def identity(t):
 
 def get_expand_reduce_stream_functions(num_streams, disable = False):
 
-    if disable:
-        return (identity, identity)
+    if num_streams == 1 or disable:
+        return (nn.Identity(), nn.Identity())
 
-    expand_fn = partial(repeat, pattern = 'b ... -> (b s) ...', s = num_streams)
-    reduce_fn = partial(reduce, pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
+    expand_fn = Reduce(pattern = 'b ... -> (b s) ...', reduction = 'repeat', s = num_streams)
+    reduce_fn = Reduce(pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
 
     return expand_fn, reduce_fn
 

hyper_connections/hyper_connections_with_multi_branch_inputs.py

@@ -11,6 +11,7 @@ from torch.nn import Module, ModuleList
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
+from einops.layers.torch import Reduce
 
 """
 ein notation:
@@ -41,11 +42,11 @@ def identity(t):
 # main functions
 
 def get_expand_reduce_stream_functions(cls, num_streams, disable = False):
-    if disable:
-        return (identity, identity)
+    if num_streams == 1 or disable:
+        return (nn.Identity(), nn.Identity())
 
-    expand_fn = partial(repeat, pattern = 'b ... -> (b s) ...', s = num_streams)
-    reduce_fn = partial(reduce, pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
+    expand_fn = Reduce(pattern = 'b ... -> (b s) ...', reduction = 'repeat', s = num_streams)
+    reduce_fn = Reduce(pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
 
     return expand_fn, reduce_fn
 

hyper_connections/hyper_connections_with_multi_input_streams.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Callable
+from typing import Callable, Union
 
 from functools import partial
 from random import randrange
@@ -11,7 +11,7 @@ from torch.nn import Module, ModuleList
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
-from einops.layers.torch import Rearrange
+from einops.layers.torch import Rearrange, Reduce
 
 """
 ein notation:
@@ -29,18 +29,15 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
-def identity(t):
-    return t
-
 # main functions
 
 def get_expand_reduce_stream_functions(num_streams, disable = False):
 
-    if disable:
-        return (identity, identity)
+    if num_streams == 1 or disable:
+        return (nn.Identity(), nn.Identity())
 
-    expand_fn = partial(repeat, pattern = 'b ... -> (b s) ...', s = num_streams)
-    reduce_fn = partial(reduce, pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
+    expand_fn = Reduce(pattern = 'b ... -> (b s) ...', reduction = 'repeat', s = num_streams)
+    reduce_fn = Reduce(pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
 
     return expand_fn, reduce_fn
 
@@ -142,7 +139,7 @@ class Residual(Module):
 
 # hyper connection with multiple input streams
 
-InputPathType = int | str  # the path to the second residual stream, where `int` points to *args[`int`] and `str` points to **kwargs[`str`] - `int` needs to be > 0, as 0 is the default input residual stream
+InputPathType = Union[int, str]  # the path to the second residual stream, where `int` points to *args[`int`] and `str` points to **kwargs[`str`] - `int` needs to be > 0, as 0 is the default input residual stream
 
 class HyperConnections(Module):
     def __init__(

{hyper_connections-0.1.4.dist-info → hyper_connections-0.1.6.dist-info}/METADATA

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

hyper_connections-0.1.6.dist-info/RECORD

@@ -0,0 +1,8 @@
+hyper_connections/__init__.py,sha256=d2zNTka0Gp9vINu4U-RhgTJFBhsVrs1fne_15Zl0oOs,224
+hyper_connections/hyper_connections.py,sha256=KBf5xiblM53jbKp_IwAgWAKQetLYWKmK8-6qNn2fHoU,10007
+hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=2JABz1slrF5_XP33L1CMNTmmixfoD464FtQpUADqneU,7806
+hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=UIKceEZEBLKFL5VuP5tR1KTDeZNIJEKjFuPAwXkcp0I,11282
+hyper_connections-0.1.6.dist-info/METADATA,sha256=3KIWntaZvm8k5BCCs_L1lOmmpB03jgEyVy8y9kYAZxI,5230
+hyper_connections-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.1.6.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.1.6.dist-info/RECORD,,

hyper_connections-0.1.4.dist-info/RECORD

@@ -1,8 +0,0 @@
-hyper_connections/__init__.py,sha256=d2zNTka0Gp9vINu4U-RhgTJFBhsVrs1fne_15Zl0oOs,224
-hyper_connections/hyper_connections.py,sha256=uOSkFjpPECFiI_WUAXMxnsvYOYlEQXo_olzIL1Cc3sk,9934
-hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=N89RSjkTbaVs46XDfPv4F770PF0kCzkMcpEuxaT5A5M,7733
-hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=_IUkvbLdv6sj6Is2ayZ7APsMzW1DZS2mjAUznKAAVLU,11258
-hyper_connections-0.1.4.dist-info/METADATA,sha256=Yd38TnztrOxNHROCAJsxNFp1m7ogRkIHZuZ_MMjHPQY,5230
-hyper_connections-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.1.4.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.1.4.dist-info/RECORD,,