fix: remove a forgotten astype

patrick-kidger · Sep 7, 2024 · 419b298 · 419b298
1 parent 5f234f6
commit 419b298
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/equinox/nn/_embedding.py b/equinox/nn/_embedding.py
@@ -185,11 +185,14 @@ def precompute_freqs_cis(
     ) -> tuple[Float[Array, "end half_emb_size"], Float[Array, "end half_emb_size"]]:
         freqs = 1.0 / (
             theta
-            ** (jnp.arange(0.0, embedding_size, 2)[jnp.newaxis, :] / embedding_size)
+            ** (
+                jnp.arange(0.0, embedding_size, 2, dtype=dtype)[jnp.newaxis, :]
+                / embedding_size
+            )
         )
 
-        t = jnp.arange(float(end))
-        freqs_outer = jnp.outer(t, freqs).astype(dtype)
+        t = jnp.arange(float(end), dtype=dtype)
+        freqs_outer = jnp.outer(t, freqs)
 
         return jnp.cos(freqs_outer), jnp.sin(freqs_outer)