hyper-connections 0.0.21__tar.gz → 0.0.22__tar.gz

{hyper_connections-0.0.21 → hyper_connections-0.0.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.21
+Version: 0.0.22
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections
@@ -34,6 +34,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
+Requires-Dist: beartype
 Requires-Dist: einops>=0.8.0
 Requires-Dist: torch>=2.3
 Provides-Extra: examples
@@ -71,9 +72,9 @@ residual = branch(residual) + residual
 
 # after, say 4 streams in paper
 
-from hyper_connections import HyperConnections
+from hyper_connections import get_init_and_expand_reduce_stream_functions
 
-init_hyper_conn, expand_stream, reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(4)
+init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(4)
 
 # 1. wrap your branch function
 
@@ -110,9 +111,9 @@ residual = branch(residual) + residual
 
 # after, say 4 streams in paper
 
-from hyper_connections import HyperConnections
+from hyper_connections import get_init_and_expand_reduce_stream_functions
 
-init_hyper_conn, expand_stream, reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(4)
+init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(4)
 
 # 1. instantiate hyper connection with correct number of streams (4 in this case) - or use the init function above
 
@@ -140,7 +141,7 @@ residual = reduce_stream(residual)
 To compare hyper connections to plain residual without changing the code, just pass `disable = True` when fetching the functions
 
 ```python
-HyperConnections.get_init_and_expand_reduce_stream_functions(4, disable = True)
+get_init_and_expand_reduce_stream_functions(4, disable = True)
 ```
 
 ## Citation

{hyper_connections-0.0.21 → hyper_connections-0.0.22}/README.md

@@ -28,9 +28,9 @@ residual = branch(residual) + residual
 
 # after, say 4 streams in paper
 
-from hyper_connections import HyperConnections
+from hyper_connections import get_init_and_expand_reduce_stream_functions
 
-init_hyper_conn, expand_stream, reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(4)
+init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(4)
 
 # 1. wrap your branch function
 
@@ -67,9 +67,9 @@ residual = branch(residual) + residual
 
 # after, say 4 streams in paper
 
-from hyper_connections import HyperConnections
+from hyper_connections import get_init_and_expand_reduce_stream_functions
 
-init_hyper_conn, expand_stream, reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(4)
+init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(4)
 
 # 1. instantiate hyper connection with correct number of streams (4 in this case) - or use the init function above
 
@@ -97,7 +97,7 @@ residual = reduce_stream(residual)
 To compare hyper connections to plain residual without changing the code, just pass `disable = True` when fetching the functions
 
 ```python
-HyperConnections.get_init_and_expand_reduce_stream_functions(4, disable = True)
+get_init_and_expand_reduce_stream_functions(4, disable = True)
 ```
 
 ## Citation

{hyper_connections-0.0.21 → hyper_connections-0.0.22}/hyper_connections/__init__.py

@@ -1,5 +1,7 @@
 from hyper_connections.hyper_connections import (
     HyperConnections,
+    get_expand_reduce_stream_functions,
+    get_init_and_expand_reduce_stream_functions,
     Residual,
     StreamEmbed,
     AttentionPoolReduceStream

{hyper_connections-0.0.21 → hyper_connections-0.0.22}/hyper_connections/hyper_connections.py

@@ -12,6 +12,8 @@ from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
 
+from beartype import beartype
+
 """
 ein notation:
 b - batch
@@ -31,6 +33,27 @@ def default(v, d):
 def identity(t):
     return t
 
+# main functions
+
+def get_expand_reduce_stream_functions(num_streams, disable = False):
+
+    if disable:
+        return (identity, 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)
+
+    return expand_fn, reduce_fn
+
+def get_init_and_expand_reduce_stream_functions(num_streams, disable = False):
+
+    hyper_conn_klass = HyperConnections if not disable else Residual
+
+    init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
+    expand_reduce_fns = get_expand_reduce_stream_functions(num_streams, disable = disable)
+
+    return (init_hyper_conn_fn, *expand_reduce_fns)
+
 # norms
 
 class RMSNorm(Module):
@@ -47,10 +70,11 @@ class RMSNorm(Module):
 # residual base class
 
 class Residual(Module):
+    @beartype
     def __init__(
         self,
         *args,
-        branch = None,
+        branch: Module | None = None,
         **kwargs
     ):
         super().__init__()
@@ -97,6 +121,7 @@ class Residual(Module):
 # hyper connection residual streams
 
 class HyperConnections(Module):
+    @beartype
     def __init__(
         self,
         num_residual_streams,
@@ -146,27 +171,6 @@ class HyperConnections(Module):
 
         self.channel_first = channel_first
 
-    @classmethod
-    def get_expand_reduce_stream_functions(cls, num_streams, disable = False):
-
-        if disable:
-            return (identity, 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)
-
-        return expand_fn, reduce_fn
-
-    @classmethod
-    def get_init_and_expand_reduce_stream_functions(cls, num_streams, disable = False):
-
-        hyper_conn_klass = cls if not disable else Residual
-
-        init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
-        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams, disable = disable)
-
-        return (init_hyper_conn_fn, *expand_reduce_fns)
-
     def width_connection(self, residuals):
         # width connection
 
@@ -244,6 +248,9 @@ class HyperConnections(Module):
 
         return add_residual_fn(branch_output)
 
+HyperConnections.get_expand_reduce_stream_functions = staticmethod(get_expand_reduce_stream_functions)
+HyperConnections.get_init_and_expand_reduce_stream_functions = staticmethod(get_init_and_expand_reduce_stream_functions)
+
 # stream embed
 
 class StreamEmbed(Module):

{hyper_connections-0.0.21 → hyper_connections-0.0.22}/hyper_connections/hyper_connections_with_multi_branch_inputs.py

@@ -12,6 +12,8 @@ from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
 
+from beartype import beartype
+
 """
 ein notation:
 b - batch
@@ -38,11 +40,32 @@ def divisible_by(num, den):
 def identity(t):
     return t
 
+# main functions
+
+def get_expand_reduce_stream_functions(cls, num_streams, disable = False):
+    if disable:
+        return (identity, 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)
+
+    return expand_fn, reduce_fn
+
+def get_init_and_expand_reduce_stream_functions(cls, num_streams, disable = False):
+
+    hyper_conn_klass = HyperConnections if not disable else Residual
+
+    init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
+    expand_reduce_fns = get_expand_reduce_stream_functions(num_streams, disable = disable)
+
+    return (init_hyper_conn_fn, *expand_reduce_fns)
+
 # main classes
 
 # hyper connection residual streams
 
 class HyperConnections(Module):
+    @beartype
     def __init__(
         self,
         num_residual_streams,
@@ -108,26 +131,6 @@ class HyperConnections(Module):
 
         self.channel_first = channel_first
 
-    @classmethod
-    def get_expand_reduce_stream_functions(cls, num_streams, disable = False):
-        if disable:
-            return (identity, 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)
-
-        return expand_fn, reduce_fn
-
-    @classmethod
-    def get_init_and_expand_reduce_stream_functions(cls, num_streams, disable = False):
-
-        hyper_conn_klass = cls if not disable else Residual
-
-        init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
-        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams, disable = disable)
-
-        return (init_hyper_conn_fn, *expand_reduce_fns)
-
     def width_connection(self, residuals):
         num_streams, num_branch_inputs = self.num_residual_streams, self.num_branch_inputs
 
@@ -225,3 +228,6 @@ class HyperConnections(Module):
         branch_output = torch.cat(branch_outputs)
 
         return add_residual_fn(branch_output)
+
+HyperConnections.get_expand_reduce_stream_functions = staticmethod(get_expand_reduce_stream_functions)
+HyperConnections.get_init_and_expand_reduce_stream_functions = staticmethod(get_init_and_expand_reduce_stream_functions)

{hyper_connections-0.0.21 → hyper_connections-0.0.22}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.0.21"
+version = "0.0.22"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -23,6 +23,7 @@ classifiers=[
 ]
 
 dependencies = [
+    "beartype",
     "einops>=0.8.0",
     "torch>=2.3",
 ]