jaxonlayers 0.2.0__tar.gz → 0.2.2__tar.gz

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

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

{jaxonlayers-0.2.0 → jaxonlayers-0.2.2}/jaxonlayers/layers/sequential.py

@@ -56,7 +56,7 @@ class BatchedLinear(eqx.Module):
         self.use_bias = use_bias
 
     def __call__(
-        self, x: Float[Array, "*batch in_features"]
+        self, x: Float[Array, "*batch in_features"], key=None
     ) -> Float[Array, "*batch out_features"]:
         input_shape = x.shape
 

{jaxonlayers-0.2.0 → jaxonlayers-0.2.2}/jaxonlayers/layers/transformer.py

@@ -40,7 +40,7 @@ class TransformerEncoderLayer(eqx.Module):
     def __init__(
         self,
         d_model: int,
-        nhead: int,
+        n_heads: int,
         dim_feedforward: int = 2048,
         dropout_p: float = 0.1,
         activation: Callable = jax.nn.relu,
@@ -50,7 +50,7 @@ class TransformerEncoderLayer(eqx.Module):
         inference: bool = False,
         *,
         key: PRNGKeyArray,
-        dtype: Any,
+        dtype: Any = None,
     ):
         if dtype is None:
             dtype = default_floating_dtype()
@@ -58,7 +58,7 @@ class TransformerEncoderLayer(eqx.Module):
         self.inference = inference
         mha_key, lin1_key, lin2_key = jax.random.split(key, 3)
         self.self_attn = eqx.nn.MultiheadAttention(
-            nhead,
+            n_heads,
             d_model,
             dropout_p=dropout_p,
             use_query_bias=use_bias,
@@ -208,7 +208,7 @@ class TransformerDecoderLayer(eqx.Module):
     def __init__(
         self,
         d_model: int,
-        nhead: int,
+        n_heads: int,
         dim_feedforward: int = 2048,
         dropout_p: float = 0.1,
         activation: Callable = jax.nn.relu,
@@ -218,7 +218,7 @@ class TransformerDecoderLayer(eqx.Module):
         inference: bool = False,
         *,
         key: PRNGKeyArray,
-        dtype: Any,
+        dtype: Any = None,
     ):
         if dtype is None:
             dtype = default_floating_dtype()
@@ -227,7 +227,7 @@ class TransformerDecoderLayer(eqx.Module):
 
         mha_key1, mha_key2, lin1_key, lin2_key = jax.random.split(key, 4)
         self.self_attn = eqx.nn.MultiheadAttention(
-            nhead,
+            n_heads,
             d_model,
             dropout_p=dropout_p,
             use_query_bias=use_bias,
@@ -239,7 +239,7 @@ class TransformerDecoderLayer(eqx.Module):
             dtype=dtype,
         )
         self.multihead_attn = eqx.nn.MultiheadAttention(
-            nhead,
+            n_heads,
             d_model,
             dropout_p=dropout_p,
             use_query_bias=use_bias,
@@ -455,7 +455,7 @@ class TransformerEncoder(eqx.Module):
     def __init__(
         self,
         d_model: int,
-        nhead: int,
+        n_heads: int,
         num_layers: int = 6,
         dim_feedforward: int = 2048,
         dropout_p: float = 0.1,
@@ -467,7 +467,7 @@ class TransformerEncoder(eqx.Module):
         inference: bool = False,
         *,
         key: PRNGKeyArray,
-        dtype: Any,
+        dtype: Any = None,
     ):
         if dtype is None:
             dtype = default_floating_dtype()
@@ -478,7 +478,7 @@ class TransformerEncoder(eqx.Module):
         self.layers = [
             TransformerEncoderLayer(
                 d_model=d_model,
-                nhead=nhead,
+                n_heads=n_heads,
                 dim_feedforward=dim_feedforward,
                 dropout_p=dropout_p,
                 activation=activation,
@@ -534,7 +534,7 @@ class TransformerDecoder(eqx.Module):
     def __init__(
         self,
         d_model: int,
-        nhead: int,
+        n_heads: int,
         num_layers: int = 6,
         dim_feedforward: int = 2048,
         dropout_p: float = 0.1,
@@ -546,7 +546,7 @@ class TransformerDecoder(eqx.Module):
         inference: bool = False,
         *,
         key: PRNGKeyArray,
-        dtype: Any,
+        dtype: Any = None,
     ):
         if dtype is None:
             dtype = default_floating_dtype()
@@ -557,7 +557,7 @@ class TransformerDecoder(eqx.Module):
         self.layers = [
             TransformerDecoderLayer(
                 d_model=d_model,
-                nhead=nhead,
+                n_heads=n_heads,
                 dim_feedforward=dim_feedforward,
                 dropout_p=dropout_p,
                 activation=activation,
@@ -627,7 +627,7 @@ class Transformer(eqx.Module):
     def __init__(
         self,
         d_model: int,
-        nhead: int,
+        n_heads: int,
         num_encoder_layers: int = 6,
         num_decoder_layers: int = 6,
         dim_feedforward: int = 2048,
@@ -639,7 +639,7 @@ class Transformer(eqx.Module):
         inference: bool = False,
         *,
         key: PRNGKeyArray,
-        dtype: Any,
+        dtype: Any = None,
     ):
         if dtype is None:
             dtype = default_floating_dtype()
@@ -650,7 +650,7 @@ class Transformer(eqx.Module):
 
         self.encoder = TransformerEncoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_encoder_layers,
             dim_feedforward=dim_feedforward,
             dropout_p=dropout_p,
@@ -666,7 +666,7 @@ class Transformer(eqx.Module):
 
         self.decoder = TransformerDecoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_decoder_layers,
             dim_feedforward=dim_feedforward,
             dropout_p=dropout_p,

{jaxonlayers-0.2.0 → jaxonlayers-0.2.2}/pyproject.toml

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

{jaxonlayers-0.2.0 → jaxonlayers-0.2.2}/tests/test_transformer.py

@@ -25,12 +25,12 @@ class TestTransformerEncoderLayer:
     )
     def test_masking(self, is_causal, use_explicit_mask):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         seq_len = 10
 
         layer = TransformerEncoderLayer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )
@@ -43,12 +43,12 @@ class TestTransformerEncoderLayer:
 
     def test_jit_no_retrace(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         seq_len = 10
 
         layer = TransformerEncoderLayer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )
@@ -84,13 +84,13 @@ class TestTransformerDecoderLayer:
         self, tgt_is_causal, memory_is_causal, use_tgt_mask, use_memory_mask
     ):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         tgt_len = 10
         src_len = 12
 
         layer = TransformerDecoderLayer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )
@@ -116,13 +116,13 @@ class TestTransformerDecoderLayer:
 
     def test_with_process_heads(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         tgt_len = 10
         src_len = 12
 
         layer = TransformerDecoderLayer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )
@@ -144,13 +144,13 @@ class TestTransformerDecoderLayer:
 
     def test_jit_no_retrace(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         tgt_len = 10
         src_len = 12
 
         layer = TransformerDecoderLayer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )
@@ -184,13 +184,13 @@ class TestTransformerEncoder:
     )
     def test_masking(self, is_causal, use_explicit_mask):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         num_layers = 3
         seq_len = 10
 
         encoder = TransformerEncoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_layers,
             key=jax.random.key(0),
             dtype=jnp.float32,
@@ -204,13 +204,13 @@ class TestTransformerEncoder:
 
     def test_with_process_heads(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         num_layers = 3
         seq_len = 10
 
         encoder = TransformerEncoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_layers,
             key=jax.random.key(0),
             dtype=jnp.float32,
@@ -226,13 +226,13 @@ class TestTransformerEncoder:
 
     def test_jit_no_retrace(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         num_layers = 3
         seq_len = 10
 
         encoder = TransformerEncoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_layers,
             key=jax.random.key(0),
             dtype=jnp.float32,
@@ -269,14 +269,14 @@ class TestTransformerDecoder:
         self, tgt_is_causal, memory_is_causal, use_tgt_mask, use_memory_mask
     ):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         num_layers = 3
         tgt_len = 10
         src_len = 12
 
         decoder = TransformerDecoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_layers,
             key=jax.random.key(0),
             dtype=jnp.float32,
@@ -303,14 +303,14 @@ class TestTransformerDecoder:
 
     def test_with_process_heads(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         num_layers = 3
         tgt_len = 10
         src_len = 12
 
         decoder = TransformerDecoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_layers,
             key=jax.random.key(0),
             dtype=jnp.float32,
@@ -333,14 +333,14 @@ class TestTransformerDecoder:
 
     def test_jit_no_retrace(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         num_layers = 3
         tgt_len = 10
         src_len = 12
 
         decoder = TransformerDecoder(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             num_layers=num_layers,
             key=jax.random.key(0),
             dtype=jnp.float32,
@@ -383,13 +383,13 @@ class TestTransformer:
         use_memory_mask,
     ):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         src_len = 12
         tgt_len = 10
 
         transformer = Transformer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )
@@ -418,13 +418,13 @@ class TestTransformer:
 
     def test_with_process_heads(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         src_len = 12
         tgt_len = 10
 
         transformer = Transformer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )
@@ -455,13 +455,13 @@ class TestTransformer:
     )
     def test_activations(self, activation):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         src_len = 12
         tgt_len = 10
 
         transformer = Transformer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             activation=activation,
             key=jax.random.key(0),
             dtype=jnp.float32,
@@ -475,13 +475,13 @@ class TestTransformer:
 
     def test_jit_no_retrace(self):
         d_model = 64
-        nhead = 4
+        n_heads = 4
         src_len = 12
         tgt_len = 10
 
         transformer = Transformer(
             d_model=d_model,
-            nhead=nhead,
+            n_heads=n_heads,
             key=jax.random.key(0),
             dtype=jnp.float32,
         )