jaxonlayers 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

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/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.dist-info → jaxonlayers-0.2.2.dist-info}/METADATA

@@ -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.dist-info → jaxonlayers-0.2.2.dist-info}/RECORD

@@ -14,9 +14,9 @@ jaxonlayers/layers/attention.py,sha256=RgtpzBPxJ4tDcUjiq_Wh_7GJndmBY6UKtbEuHGLA1
 jaxonlayers/layers/convolution.py,sha256=k0dMFBDjzycB7UNuyHqKihJtBa6u93V6OLxyUUyipN4,3247
 jaxonlayers/layers/normalization.py,sha256=3aGzNzDN05A72ZHLUM2w9WpicLtGsjzj1l0jhuyn63U,8379
 jaxonlayers/layers/regularization.py,sha256=ZrvtBJPH84xuxrxEbZc7TBxjp8OvKEv4ecan5s8F9zs,563
-jaxonlayers/layers/sequential.py,sha256=Tw98hNZiXMC-CYZD6h_pi7eAxkgHeQAUvZF2I9H0d8Y,2833
+jaxonlayers/layers/sequential.py,sha256=xBZavhSra0oZUZjSlThzIGER5xR62n46mgTkPqRV2Y0,2843
 jaxonlayers/layers/state_space.py,sha256=oDVRbduNtU48Q4rLd-XywZcqVN0QYTlq1UUhOXcGLoo,2537
-jaxonlayers/layers/transformer.py,sha256=u2nRLFp1cnM2WvAXQQ5zKyHbLFV0bUnaW91x56dl6n8,21781
-jaxonlayers-0.2.0.dist-info/METADATA,sha256=gTtPTaJzmBNNMZcTbe-gBVFT09GPjCGVUbkRVv_7hhc,565
-jaxonlayers-0.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-jaxonlayers-0.2.0.dist-info/RECORD,,
+jaxonlayers/layers/transformer.py,sha256=Syuh_kN-I-Gg-C20trQcNio1WEG88OHDRl3ShFaQNqQ,21848
+jaxonlayers-0.2.2.dist-info/METADATA,sha256=sYqihkNlwRrIMMPtmi8zIMPk4TOef3H6ZHB44gqKINM,565
+jaxonlayers-0.2.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+jaxonlayers-0.2.2.dist-info/RECORD,,