jaxonlayers 0.1.4__tar.gz → 0.2.0__tar.gz

{jaxonlayers-0.1.4 → jaxonlayers-0.2.0}/.pre-commit-config.yaml

@@ -1,12 +1,12 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.13.1
+    rev: v0.14.7
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.405
+    rev: v1.1.407
     hooks:
       - id: pyright
         additional_dependencies:

{jaxonlayers-0.1.4 → jaxonlayers-0.2.0}/PKG-INFO

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

{jaxonlayers-0.1.4 → jaxonlayers-0.2.0}/jaxonlayers/functions/masking.py

@@ -1,4 +1,5 @@
 import jax.numpy as jnp
+from jaxtyping import Array, Bool
 
 
 def canonical_mask(
@@ -50,7 +51,27 @@ def canonical_attn_mask(attn_mask, query_dtype=jnp.float32):
     )
 
 
+def make_causal_mask(seq_len: int) -> Bool[Array, "seq_len seq_len"]:
+    """
+    Returns a boolean mask.
+
+    Example:
+    [[True,  False, False],
+     [True,  True,  False],
+     [True,  True,  True ]]
+    """
+    return jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))
+
+
 def build_attention_mask(context_length: int):
+    """
+    Returns a numerical matrix with 0 and -inf.
+
+    Example:
+    [[ 0,   -inf, -inf],
+     [ 0,    0,   -inf],
+     [ 0,    0,    0  ]]
+    """
     mask = jnp.tril(jnp.zeros((context_length, context_length)))
     upper = jnp.triu(jnp.full((context_length, context_length), float("-inf")), k=1)
 

{jaxonlayers-0.1.4 → jaxonlayers-0.2.0}/jaxonlayers/functions/utils.py

@@ -5,7 +5,7 @@ from jaxtyping import PyTree
 
 
 def default_floating_dtype():
-    if jax.config.jax_enable_x64:  # pyright: ignore
+    if jax.config.read("jax_enable_x64"):  # pyright: ignore
         return jnp.float64
     else:
         return jnp.float32

{jaxonlayers-0.1.4 → jaxonlayers-0.2.0}/jaxonlayers/layers/__init__.py

@@ -5,6 +5,13 @@ from .normalization import BatchNorm, LayerNorm, LocalResponseNormalization
 from .regularization import StochasticDepth
 from .sequential import BatchedLinear
 from .state_space import SelectiveStateSpace
+from .transformer import (
+    Transformer,
+    TransformerDecoder,
+    TransformerDecoderLayer,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
 
 __all__ = [
     "BatchNorm",
@@ -18,4 +25,9 @@ __all__ = [
     "AbstractNormStateful",
     "AbstractNorm",
     "BatchedLinear",
+    "Transformer",
+    "TransformerDecoder",
+    "TransformerDecoderLayer",
+    "TransformerEncoder",
+    "TransformerEncoderLayer",
 ]

{jaxonlayers-0.1.4 → jaxonlayers-0.2.0}/jaxonlayers/layers/attention.py

@@ -131,6 +131,7 @@ class MultiheadAttention(eqx.Module):
         attn_mask: Array | None = None,
         average_attn_weights: bool = True,
         is_causal: bool = False,
+        dropout_key: PRNGKeyArray | None = None,
     ) -> tuple[Array, Array | None]:
         key_padding_mask = canonical_mask(
             mask=key_padding_mask,
@@ -171,6 +172,7 @@ class MultiheadAttention(eqx.Module):
                 v_proj_weight=self.v_proj_weight,
                 average_attn_weights=average_attn_weights,
                 is_causal=is_causal,
+                dropout_key=dropout_key,
             )
         else:
             attn_output, attn_output_weights = multi_head_attention_forward(
@@ -193,6 +195,7 @@ class MultiheadAttention(eqx.Module):
                 attn_mask=attn_mask,
                 average_attn_weights=average_attn_weights,
                 is_causal=is_causal,
+                dropout_key=dropout_key,
             )
 
         return attn_output, attn_output_weights