hyper-connections 0.1.4__tar.gz → 0.1.6__tar.gz

{hyper_connections-0.1.4 → hyper_connections-0.1.6}/PKG-INFO

@@ -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.4 → hyper_connections-0.1.6}/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-0.1.4 → hyper_connections-0.1.6}/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-0.1.4 → hyper_connections-0.1.6}/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 → hyper_connections-0.1.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.1.4"
+version = "0.1.6"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }