hyper-connections 0.3.6__py3-none-any.whl → 0.3.7__py3-none-any.whl

hyper_connections/manifold_constrained_hyper_connections.py

@@ -46,6 +46,9 @@ def l1norm(t, dim):
     return F.normalize(t, p = 1, dim = dim)
 
 def sinkhorn_knopps(log_alpha, iters = 20):
+    dtype = log_alpha.dtype
+    log_alpha = log_alpha.float()
+
     log_alpha = log_alpha - log_alpha.amax(dim = -2, keepdim = True).detach()
 
     alpha = log_alpha.exp()
@@ -54,7 +57,7 @@ def sinkhorn_knopps(log_alpha, iters = 20):
         alpha = l1norm(alpha, dim = -2)
         alpha = l1norm(alpha, dim = -1)
 
-    return alpha
+    return alpha.to(dtype)
 
 # main functions
 
@@ -316,17 +319,21 @@ class ManifoldConstrainedHyperConnections(Module):
 
         # alpha for weighted sum of residuals going into branch
 
-        wc_weight = normed @ self.dynamic_alpha_fn
+        dtype = residuals.dtype
+
+        normed = normed.float()
+
+        wc_weight = normed @ self.dynamic_alpha_fn.float()
 
-        pre_branch_scale = repeat(self.pre_branch_scale, '1 -> s', s = self.num_fracs)
-        residual_scale = repeat(self.residual_scale, '1 -> s', s = self.num_fracs * streams)
+        pre_branch_scale = repeat(self.pre_branch_scale.float(), '1 -> s', s = self.num_fracs)
+        residual_scale = repeat(self.residual_scale.float(), '1 -> s', s = self.num_fracs * streams)
         alpha_scale = cat((pre_branch_scale, residual_scale))
 
         alpha_scale = repeat(alpha_scale, 'n -> (v n)', v = self.num_input_views)
 
         dynamic_alpha = wc_weight * alpha_scale
 
-        static_alpha = rearrange(self.static_alpha, '(f s) d -> f s d', s = streams)
+        static_alpha = rearrange(self.static_alpha.float(), '(f s) d -> f s d', s = streams)
 
         alpha = dynamic_alpha + static_alpha
 
@@ -351,20 +358,20 @@ class ManifoldConstrainedHyperConnections(Module):
         beta = None
 
         if self.add_branch_out_to_residual:
-            dc_weight = normed @ self.dynamic_beta_fn
+            dc_weight = normed @ self.dynamic_beta_fn.float()
 
             dc_weight = dc_weight.sigmoid() * 2 # sigmoid * 2 for "H_post", corresponding to dc weight in original paper
 
             if not self.has_fracs:
                 dc_weight = rearrange(dc_weight, '... -> ... 1')
 
-            dynamic_beta = dc_weight * self.h_post_scale
+            dynamic_beta = dc_weight * self.h_post_scale.float()
 
-            static_beta = rearrange(self.static_beta, '... (s f) -> ... s f', s = streams)
+            static_beta = rearrange(self.static_beta.float(), '... (s f) -> ... s f', s = streams)
 
             beta = dynamic_beta + static_beta
 
-        mix_h = einsum(alpha, residuals, '... f1 s f2 t, ... f1 s d -> ... f2 t d')
+        mix_h = einsum(alpha, residuals.float(), '... f1 s f2 t, ... f1 s d -> ... f2 t d')
 
         if self.num_input_views == 1:
             branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :]
@@ -379,6 +386,12 @@ class ManifoldConstrainedHyperConnections(Module):
 
         branch_input = self.merge_fracs(branch_input)
 
+        branch_input = branch_input.to(dtype)
+        residuals = residuals.to(dtype)
+
+        if exists(beta):
+            beta = beta.to(dtype)
+
         return branch_input, maybe_transformed_residuals, dict(beta = beta)
 
     def depth_connection(
@@ -399,7 +412,9 @@ class ManifoldConstrainedHyperConnections(Module):
         if self.channel_first:
             branch_output = rearrange(branch_output, 'b d ... -> b ... d')
 
-        output = einsum(branch_output, beta, 'b ... f1 d, b ... f1 s f2 -> b ... f2 s d')
+        dtype = residuals.dtype
+
+        output = einsum(branch_output.float(), beta.float(), 'b ... f1 d, b ... f1 s f2 -> b ... f2 s d')
 
         output = rearrange(output, 'b ... s d -> (b s) ... d')
 
@@ -412,7 +427,7 @@ class ManifoldConstrainedHyperConnections(Module):
         if self.channel_first:
             output = rearrange(output, 'b ... d -> b d ...')
 
-        residuals = self.depth_residual_fn(output, residuals)
+        residuals = self.depth_residual_fn(output.to(dtype), residuals)
 
         return self.dropout(residuals)
 

{hyper_connections-0.3.6.dist-info → hyper_connections-0.3.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.3.6
+Version: 0.3.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.3.6.dist-info → hyper_connections-0.3.7.dist-info}/RECORD

@@ -3,9 +3,9 @@ hyper_connections/hyper_connections.py,sha256=UHxZhyRwx89GRgmQVt53Gv6JeNhX8UCjjE
 hyper_connections/hyper_connections_channel_first.py,sha256=_1PM4LRcPpDqfCiHlBMc2nLV08sXM2nuyZGSKTiuqbE,6818
 hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=6BXKdSwyx6wdQVseebKG2EQkhVaVLrrepOlL8lLnex4,7855
 hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=ueT3CJPHrt5hRU7q1bFF0rANWJh_pXqclt6HiUu1gBY,11331
-hyper_connections/manifold_constrained_hyper_connections.py,sha256=ruPwN4mo2Vo5e2KAIQEeeRrJr47Ds1mapM5DfqrwTq8,16552
+hyper_connections/manifold_constrained_hyper_connections.py,sha256=gJzc9oHZjhC3S85HiXGQck5qBNcRWGCjVsDfMXwqPxo,16961
 hyper_connections/residuals.py,sha256=JVSFJj_H7xQ3_Fd-pZH5Hdv9SveAQu29jQNvMyom5ek,921
-hyper_connections-0.3.6.dist-info/METADATA,sha256=rpv7M5wjGiqo1q5AlaMXmfmDDqJ6gR-ekftkOxqXbDk,6704
-hyper_connections-0.3.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hyper_connections-0.3.6.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.3.6.dist-info/RECORD,,
+hyper_connections-0.3.7.dist-info/METADATA,sha256=Pyj7qVEMj6Szb_PD80eYoN3O4fwv8DjsYSzFd8EY_bo,6704
+hyper_connections-0.3.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hyper_connections-0.3.7.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.3.7.dist-info/RECORD,,