jaxonlayers 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

jaxonlayers/functions/activation.py

@@ -5,4 +5,4 @@ from jaxtyping import Array
 
 def swiglu(x: Array, axis=-1):
     a, b = jnp.split(x, 2, axis=axis)
-    return a * jax.nn.swish(b)
+    return jax.nn.swish(a) * b

jaxonlayers/layers/embedding.py

@@ -0,0 +1,56 @@
+import equinox as eqx
+from beartype.typing import Any
+from jaxtyping import Array, PRNGKeyArray
+
+from jaxonlayers.functions import default_floating_dtype
+
+
+class EmbeddingWithPadding(eqx.Module):
+    embed: eqx.nn.Embedding
+    padding_idx: int = eqx.field(static=True)
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        padding_idx: int = 0,
+        *,
+        key: PRNGKeyArray,
+        dtype: Any | None = None,
+    ):
+        if dtype is None:
+            dtype = default_floating_dtype()
+        assert dtype is not None
+        self.embed = eqx.nn.Embedding(
+            num_embeddings, embedding_dim, key=key, dtype=dtype
+        )
+        self.padding_idx = padding_idx
+
+    def __call__(self, x: Array):
+        out = self.embed(x)
+        mask = (x != self.padding_idx).astype(out.dtype)
+        return out * mask[..., None]
+
+
+class EmbeddingBag(eqx.Module):
+    embed: EmbeddingWithPadding
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        padding_idx: int = 0,
+        *,
+        key: PRNGKeyArray,
+        dtype: Any | None = None,
+    ):
+        if dtype is None:
+            dtype = default_floating_dtype()
+        assert dtype is not None
+        self.embed = EmbeddingWithPadding(
+            num_embeddings, embedding_dim, padding_idx=padding_idx, key=key, dtype=dtype
+        )
+
+    def __call__(self, x):
+        looked_up = eqx.filter_vmap(self.embed)(x)
+        return looked_up.sum(axis=0)

{jaxonlayers-0.2.3.dist-info → jaxonlayers-0.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxonlayers
-Version: 0.2.3
+Version: 0.2.5
 Summary: Additional layers and functions that extend Equinox
 Requires-Python: >=3.13
 Requires-Dist: beartype>=0.21.0

{jaxonlayers-0.2.3.dist-info → jaxonlayers-0.2.5.dist-info}/RECORD

@@ -1,6 +1,6 @@
 jaxonlayers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxonlayers/functions/__init__.py,sha256=lxMNSnEn2pJ1XLLyQTTFXiCcJTdrzfE1RlorqD9lEog,711
-jaxonlayers/functions/activation.py,sha256=cjwDEadbJECt1AE1uEtKV_p_wBK_pJ7Bdud0ebm_GNQ,165
+jaxonlayers/functions/activation.py,sha256=cVR1Obcv5JE7GiUNYlxHCcf5bpqZ9OQq8zYwSBJ-svY,165
 jaxonlayers/functions/attention.py,sha256=AyaAeA2yo5Cgljk2rU6JlVttQeG0FMBLE-f7285PdM0,12945
 jaxonlayers/functions/embedding.py,sha256=GsXxpB3G72aLY1SHdZk7gxKbNftlEZ3HRDy-7_5z4HM,610
 jaxonlayers/functions/initialization.py,sha256=h7uzdPl-rL7faT9hbRs5aN6EeaX0r70y88ad1cwZvmY,1161
@@ -13,11 +13,12 @@ jaxonlayers/layers/__init__.py,sha256=aSYMojd0wKdGLxrPxZi_x1c4n8JTJvjAr2vNXC1zYP
 jaxonlayers/layers/abstract.py,sha256=uyPKGsH5DETnjiU3PUiF98tk9boNZW96amxrM5JQlZY,366
 jaxonlayers/layers/attention.py,sha256=RgtpzBPxJ4tDcUjiq_Wh_7GJndmBY6UKtbEuHGLA1PY,7915
 jaxonlayers/layers/convolution.py,sha256=k0dMFBDjzycB7UNuyHqKihJtBa6u93V6OLxyUUyipN4,3247
+jaxonlayers/layers/embedding.py,sha256=G0LyBZpKee4S6HqhgxdH6rCJNZwMomKNeDmYqRMIPGI,1509
 jaxonlayers/layers/normalization.py,sha256=3aGzNzDN05A72ZHLUM2w9WpicLtGsjzj1l0jhuyn63U,8379
 jaxonlayers/layers/regularization.py,sha256=ZrvtBJPH84xuxrxEbZc7TBxjp8OvKEv4ecan5s8F9zs,563
 jaxonlayers/layers/sequential.py,sha256=xBZavhSra0oZUZjSlThzIGER5xR62n46mgTkPqRV2Y0,2843
 jaxonlayers/layers/state_space.py,sha256=oDVRbduNtU48Q4rLd-XywZcqVN0QYTlq1UUhOXcGLoo,2537
 jaxonlayers/layers/transformer.py,sha256=Syuh_kN-I-Gg-C20trQcNio1WEG88OHDRl3ShFaQNqQ,21848
-jaxonlayers-0.2.3.dist-info/METADATA,sha256=-Pde3UQPCm39v_dIWtAZ4FjKE8CRgl8WbVbsM5-qbAM,565
-jaxonlayers-0.2.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-jaxonlayers-0.2.3.dist-info/RECORD,,
+jaxonlayers-0.2.5.dist-info/METADATA,sha256=TaYvkI9lsjDOwKl-rXCnYXZlof9xh1AgxdANXyWzn2o,565
+jaxonlayers-0.2.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+jaxonlayers-0.2.5.dist-info/RECORD,,