jaxonlayers 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (5) hide show
  1. jaxonlayers/functions/embedding.py +1 -1
  2. jaxonlayers/layers/__init__.py +5 -0
  3. {jaxonlayers-0.1.3.dist-info → jaxonlayers-0.1.4.dist-info}/METADATA +2 -2
  4. {jaxonlayers-0.1.3.dist-info → jaxonlayers-0.1.4.dist-info}/RECORD +5 -5
  5. {jaxonlayers-0.1.3.dist-info → jaxonlayers-0.1.4.dist-info}/WHEEL +0 -0
jaxonlayers/functions/embedding.py CHANGED
@@ -16,6 +16,6 @@ def sinusoidal_embedding(
16
16
  )
17
17
 
18
18
  time_args = t * embedding_freqs
19
- embedding = jnp.concatenate([jnp.sin(time_args), jnp.cos(time_args)])
19
+ embedding = jnp.concatenate([jnp.sin(time_args), jnp.cos(time_args)], axis=-1)
20
20
 
21
21
  return embedding
jaxonlayers/layers/__init__.py CHANGED
@@ -1,7 +1,9 @@
1
+ from .abstract import AbstractNorm, AbstractNormStateful
1
2
  from .attention import MultiheadAttention, SqueezeExcitation
2
3
  from .convolution import ConvNormActivation
3
4
  from .normalization import BatchNorm, LayerNorm, LocalResponseNormalization
4
5
  from .regularization import StochasticDepth
6
+ from .sequential import BatchedLinear
5
7
  from .state_space import SelectiveStateSpace
6
8
 
7
9
  __all__ = [
@@ -13,4 +15,7 @@ __all__ = [
13
15
  "StochasticDepth",
14
16
  "ConvNormActivation",
15
17
  "LayerNorm",
18
+ "AbstractNormStateful",
19
+ "AbstractNorm",
20
+ "BatchedLinear",
16
21
  ]
{jaxonlayers-0.1.3.dist-info → jaxonlayers-0.1.4.dist-info}/METADATA RENAMED
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: jaxonlayers
3
- Version: 0.1.3
4
- Summary: Add your description here
3
+ Version: 0.1.4
4
+ Summary: Additional layers and functions that extend Equinox
5
5
  Requires-Python: >=3.13
6
6
  Requires-Dist: beartype>=0.21.0
7
7
  Requires-Dist: equinox>=0.13.0
{jaxonlayers-0.1.3.dist-info → jaxonlayers-0.1.4.dist-info}/RECORD RENAMED
@@ -1,14 +1,14 @@
1
1
  jaxonlayers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
2
2
  jaxonlayers/functions/__init__.py,sha256=lxMNSnEn2pJ1XLLyQTTFXiCcJTdrzfE1RlorqD9lEog,711
3
3
  jaxonlayers/functions/attention.py,sha256=AyaAeA2yo5Cgljk2rU6JlVttQeG0FMBLE-f7285PdM0,12945
4
- jaxonlayers/functions/embedding.py,sha256=g76Wg_MSyqatwjWrhDSfAxDaMyr2s3A5AoST_cRa3Q8,601
4
+ jaxonlayers/functions/embedding.py,sha256=GsXxpB3G72aLY1SHdZk7gxKbNftlEZ3HRDy-7_5z4HM,610
5
5
  jaxonlayers/functions/initialization.py,sha256=h7uzdPl-rL7faT9hbRs5aN6EeaX0r70y88ad1cwZvmY,1161
6
6
  jaxonlayers/functions/masking.py,sha256=5VeHMuoVeKxAHtzGb74GqvKHfib4wJAETTAn1oaOgLM,1531
7
7
  jaxonlayers/functions/normalization.py,sha256=e2vNNbDz-Y6j5bgQshd4MshnzLcDOjUnGSinpRXtPtA,206
8
8
  jaxonlayers/functions/regularization.py,sha256=ltFVeJZFhdNrHsH0CjDl9Y3dlAQmh1ABntCjpYUBgXM,1317
9
9
  jaxonlayers/functions/state_space.py,sha256=_PdykvDSRZWmIQR5KZTUN6shLhPGElio9THnoVpLJ_g,2747
10
10
  jaxonlayers/functions/utils.py,sha256=M1uc01yOADy2ig6YsIyn0JBf_zdFuJS-rxin5RVub-A,1749
11
- jaxonlayers/layers/__init__.py,sha256=gjH0QUOCplii5KNYshWNehPfHZLVjEeYk8EhhY7FHRE,480
11
+ jaxonlayers/layers/__init__.py,sha256=FXHlQXpJieDgXFy0HyYiZU8VqlwHaWbuDF8EI0jd57k,644
12
12
  jaxonlayers/layers/abstract.py,sha256=uyPKGsH5DETnjiU3PUiF98tk9boNZW96amxrM5JQlZY,366
13
13
  jaxonlayers/layers/attention.py,sha256=NdjDnA3yYFpvXD2_SjBwF-8Pdk_GeRLPuPekTR-iMwE,7784
14
14
  jaxonlayers/layers/convolution.py,sha256=k0dMFBDjzycB7UNuyHqKihJtBa6u93V6OLxyUUyipN4,3247
@@ -16,6 +16,6 @@ jaxonlayers/layers/normalization.py,sha256=3aGzNzDN05A72ZHLUM2w9WpicLtGsjzj1l0jh
16
16
  jaxonlayers/layers/regularization.py,sha256=ZrvtBJPH84xuxrxEbZc7TBxjp8OvKEv4ecan5s8F9zs,563
17
17
  jaxonlayers/layers/sequential.py,sha256=Tw98hNZiXMC-CYZD6h_pi7eAxkgHeQAUvZF2I9H0d8Y,2833
18
18
  jaxonlayers/layers/state_space.py,sha256=oDVRbduNtU48Q4rLd-XywZcqVN0QYTlq1UUhOXcGLoo,2537
19
- jaxonlayers-0.1.3.dist-info/METADATA,sha256=HHbeosdFV7GvUHihj3i6udPQSgTa01ISU8AAVEM6LYA,539
20
- jaxonlayers-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
21
- jaxonlayers-0.1.3.dist-info/RECORD,,
19
+ jaxonlayers-0.1.4.dist-info/METADATA,sha256=PxlTeg0WsEBAONPHDmgsAyIMidTRGmlPQ6ioSwbQLIg,565
20
+ jaxonlayers-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
21
+ jaxonlayers-0.1.4.dist-info/RECORD,,