jaxonlayers 0.2.4__tar.gz → 0.2.5__tar.gz

{jaxonlayers-0.2.4 → jaxonlayers-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxonlayers
-Version: 0.2.4
+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.5/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.4 → jaxonlayers-0.2.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxonlayers"
-version = "0.2.4"
+version = "0.2.5"
 description = "Additional layers and functions that extend Equinox"
 readme = "README.md"
 requires-python = ">=3.13"