hyper-connections 0.0.24__tar.gz → 0.1.1__tar.gz

hyper_connections-0.1.1/.github/workflows/test.yml

@@ -0,0 +1,19 @@
+name: Tests the examples in README
+on: push
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install Python
+        uses: actions/setup-python@v4
+      - name: Install the latest version of rye
+        uses: eifinger/setup-rye@v2
+      - name: Use UV instead of pip
+        run: rye config --set-bool behavior.use-uv=true
+      - name: Install dependencies
+        run: |
+          rye sync
+      - name: Run pytest
+        run: rye run pytest tests/

{hyper_connections-0.0.24 → hyper_connections-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.24
+Version: 0.1.1
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections
@@ -34,12 +34,9 @@ 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
-Provides-Extra: test
-Requires-Dist: pytest; extra == 'test'
 Description-Content-Type: text/markdown
 
 <img src="./hyper-connections.png" width="450px"></img>

{hyper_connections-0.0.24 → hyper_connections-0.1.1}/hyper_connections/hyper_connections_with_multi_branch_inputs.py

@@ -12,8 +12,6 @@ from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
 
-from beartype import beartype
-
 """
 ein notation:
 b - batch
@@ -65,7 +63,6 @@ def get_init_and_expand_reduce_stream_functions(cls, num_streams, disable = Fals
 # hyper connection residual streams
 
 class HyperConnections(Module):
-    @beartype
     def __init__(
         self,
         num_residual_streams,

{hyper_connections-0.0.24 → hyper_connections-0.1.1}/hyper_connections/hyper_connections_with_multi_input_streams.py

@@ -13,8 +13,6 @@ from torch.utils._pytree import tree_flatten, tree_unflatten
 from einops import rearrange, repeat, reduce, einsum
 from einops.layers.torch import Rearrange
 
-from beartype import beartype
-
 """
 ein notation:
 b - batch
@@ -95,7 +93,6 @@ class ProjActScale(Module):
 # residual base class
 
 class Residual(Module):
-    @beartype
     def __init__(
         self,
         *args,
@@ -145,10 +142,9 @@ 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` + 1] and `str` points to **kwargs[`str`]
+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
 
 class HyperConnections(Module):
-    @beartype
     def __init__(
         self,
         num_residual_streams,
@@ -185,7 +181,7 @@ class HyperConnections(Module):
         init_alpha0 = torch.zeros((num_residual_streams, 1))
         init_alpha0[init_residual_index, 0] = 1.
 
-        self.dynamic_alpha_and_branch_input = ProjActScale(dim, num_residual_streams + 1)
+        self.dynamic_alpha_and_branch_input = ProjActScale(dim, num_residual_streams + 1, activation = act)
         self.static_alpha = nn.Parameter(torch.cat([init_alpha0, torch.eye(num_residual_streams)], dim = 1))
 
         self.dynamic_beta = ProjActScale(dim, 1, activation = act, squeeze_output = True)
@@ -200,7 +196,7 @@ class HyperConnections(Module):
 
         self.additional_norms = ModuleList([RMSNorm(dim) for _, dim in additional_input_paths])
         self.additional_to_dynamic_input = ModuleList([ProjActScale(dim, 1, activation = act, squeeze_output = True) for _ , dim in additional_input_paths])
-        self.additional_static_input = nn.ParameterList([nn.Parameter(init_alpha0[..., 0])])
+        self.additional_static_input = nn.ParameterList([nn.Parameter(init_alpha0[..., 0]) for _ in additional_input_paths])
 
         self.additional_input_paths = additional_input_paths
 

{hyper_connections-0.0.24 → hyper_connections-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.0.24"
+version = "0.1.1"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -23,7 +23,6 @@ classifiers=[
 ]
 
 dependencies = [
-    "beartype",
     "einops>=0.8.0",
     "torch>=2.3",
 ]
@@ -34,14 +33,9 @@ Repository = "https://github.com/lucidrains/hyper-connections"
 
 [project.optional-dependencies]
 examples = []
-test = [
-    "pytest"
-]
 
 [tool.pytest.ini_options]
-pythonpath = [
-  "."
-]
+pythonpath = ["."]
 
 [build-system]
 requires = ["hatchling"]
@@ -49,7 +43,7 @@ build-backend = "hatchling.build"
 
 [tool.rye]
 managed = true
-dev-dependencies = []
+dev-dependencies = ["pytest>=8.2.0"]
 
 [tool.hatch.metadata]
 allow-direct-references = true

hyper_connections-0.1.1/tests/test_hyper_connections.py

@@ -0,0 +1,138 @@
+import pytest
+
+import torch
+from torch import nn
+
+@pytest.mark.parametrize('disable', (False, True))
+def test_readme(disable):
+
+    # a single branch layer
+
+    branch = nn.Linear(512, 512)
+
+    # before
+
+    residual = torch.randn(2, 1024, 512)
+
+    residual = branch(residual) + residual
+
+    # after, say 4 streams in paper
+
+    from hyper_connections import get_init_and_expand_reduce_stream_functions
+
+    init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(4, disable = disable)
+
+    # 1. wrap your branch function
+
+    hyper_conn_branch = init_hyper_conn(dim = 512, branch = branch)
+
+    # 2. expand to 4 streams, this must be done before your trunk, typically a for-loop with many branch functions
+
+    residual = expand_stream(residual)
+
+    # 3. forward your residual as usual into the wrapped branch function(s)
+
+    residual = hyper_conn_branch(residual) 
+
+    # 4. reduce 4 streams with a summation, this has to be done after your for-loop trunk. for transformer, unsure whether to do before or after final norm
+
+    residual = reduce_stream(residual)
+
+    assert residual.shape == (2, 1024, 512)
+
+def test_manual():
+    # a single branch layer
+
+    branch = nn.Linear(512, 512)
+
+    # before
+
+    residual = torch.randn(2, 1024, 512)
+
+    residual = branch(residual) + residual
+
+    # after, say 4 streams in paper
+
+    from hyper_connections import get_init_and_expand_reduce_stream_functions
+
+    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
+
+    hyper_conn = init_hyper_conn(dim = 512)
+
+    # 2. expand to 4 streams
+
+    residual = expand_stream(residual)
+
+    # 3. forward your residual into hyper connection for the branch input + add residual function (learned betas)
+
+    branch_input, add_residual = hyper_conn(residual)
+
+    branch_output = branch(branch_input)
+
+    residual = add_residual(branch_output)
+
+    # or you can do it in one line as so -> residual = hyper_conn.decorate_branch(branch)(residual)
+
+    # 4. reduce 4 streams with a summation, this has to be done after your for loop trunk
+
+    residual = reduce_stream(residual)
+    assert residual.shape == (2, 1024, 512)
+
+@pytest.mark.parametrize('disable', (False, True))
+def test_multi_input_hyper_connections(disable):
+
+    # two branch layers
+
+    class CustomModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = nn.Linear(512, 512)
+            self.second_linear = nn.Linear(256, 512)
+            self.third_linear = nn.Linear(128, 512)
+
+        def forward(self, x, second, *, third):
+            return self.linear(x) + self.second_linear(second) + self.third_linear(third), 3.
+
+    branch = CustomModule()
+
+    # before
+
+    residual = torch.randn(3, 1024, 512)
+    second_residual = torch.randn(3, 1024, 256)
+    third_residual = torch.randn(3, 1024, 128)
+
+    # residual = branch1(residual) + branch2(residual) + residual
+
+    # after, say 4 streams in paper
+
+    from hyper_connections.hyper_connections_with_multi_input_streams import HyperConnections
+
+    init_hyper_conn, expand_stream, reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(4, disable = disable)
+
+    # 1. instantiate hyper connection with correct number of streams (4 in this case) - or use the init function above
+
+    hyper_conn = init_hyper_conn(
+        dim = 512,
+        branch = branch,
+        additional_input_paths = [
+            (1, 256),        # points at second residual stream, first arg
+            ('third', 128)   # points at third residual stream, keyword argument 'third'
+        ],
+        layer_index = 1,
+    )
+
+    # 2. expand to 4 streams
+
+    residual = expand_stream(residual)
+    second_residual = expand_stream(second_residual)
+    third_residual = expand_stream(third_residual)
+
+    # 3. forward your residual into hyper connection for the branch input + add residual function (learned betas)
+
+    residual, rest_output = hyper_conn(residual, second_residual, third = third_residual)
+
+    residual = reduce_stream(residual)
+
+    assert residual.shape == (3, 1024, 512)