Norm should return float rather than complex

patrick-kidger · Mar 23, 2024 · 48eef77 · 48eef77
1 parent 1d0ac6c
commit 48eef77
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lineax/_norm.py b/lineax/_norm.py
@@ -113,7 +113,7 @@ def rms_norm(x: PyTree[ArrayLike]) -> Scalar:
         if len(leaves) == 0:
             dtype = default_floating_dtype()
         else:
-            dtype = jnp.result_type(*leaves)
+            dtype = jnp.finfo(jnp.result_type(*leaves)).dtype
         return jnp.array(0.0, dtype)
     else:
         return two_norm(x) / math.sqrt(size)
@@ -131,7 +131,7 @@ def max_norm(x: PyTree[ArrayLike]) -> Scalar:
         if len(leaves) == 0:
             dtype = default_floating_dtype()
         else:
-            dtype = jnp.result_type(*leaves)
+            dtype = jnp.finfo(jnp.result_type(*leaves)).dtype
         return jnp.array(0.0, dtype)
     else:
         out = ft.reduce(jnp.maximum, leaf_maxes)