nmn 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

nmn/nnx/squashers/softermax.py

@@ -1,7 +1,12 @@
+from functools import partial
+from typing import Optional
+
+import jax
 import jax.numpy as jnp
 from jax import Array
-from typing import Optional
 
+
+@partial(jax.jit, static_argnames=("n", "axis"))
 def softermax(
     x: Array,
     n: float = 1.0,

nmn/nnx/yatattention.py

@@ -44,6 +44,8 @@ def yat_attention_weights(
   module: Optional[Module] = None,
   epsilon: float = 1e-5,
   use_softermax: bool = False,
+  power: float = 1.0,
+
 ):
   """Computes attention weights using YatNMN distance-based calculation."""
   query, key = promote_dtype((query, key), dtype=dtype)
@@ -89,7 +91,7 @@ def yat_attention_weights(
 
   # normalize the attention weights
   if use_softermax:
-    attn_weights = softermax(attn_weights).astype(dtype)
+    attn_weights = softermax(attn_weights, n=power).astype(dtype)
   else:
     attn_weights = jax.nn.softmax(attn_weights).astype(dtype)
 
@@ -126,6 +128,7 @@ def yat_attention(
   module: Optional[Module] = None,
   epsilon: float = 1e-5,
   use_softermax: bool = False,
+  power: float = 1.0,
 ):
   """Computes attention using YatNMN distance-based calculation."""
   query, key, value = promote_dtype((query, key, value), dtype=dtype)
@@ -154,6 +157,7 @@ def yat_attention(
     module,
     epsilon,
     use_softermax,
+    power,
   )
 
   # return weighted sum over values for each query position
@@ -370,6 +374,7 @@ class MultiHeadAttention(Module):
     rngs: rnglib.Rngs,
     epsilon: float = 1e-5,
     use_softermax: bool = False,
+    power: float = 1.0,
   ):
     self.num_heads = num_heads
     self.in_features = in_features
@@ -399,6 +404,7 @@ class MultiHeadAttention(Module):
     self.out_dot_general_cls = out_dot_general_cls
     self.epsilon = epsilon
     self.use_softermax = use_softermax
+    self.power = power
     self.use_alpha = use_alpha
     self.alpha_init = alpha_init
     self.use_dropconnect = use_dropconnect
@@ -627,6 +633,7 @@ class MultiHeadAttention(Module):
       module=self if sow_weights else None,
       epsilon=self.epsilon,  # Pass epsilon to yat_attention
       use_softermax=self.use_softermax,
+      power= self.power,
     )
     # Reshape attention output back to original embedding dimension
     # from [batch..., length, num_heads, head_dim] to [batch..., length, qkv_features]

{nmn-0.1.8.dist-info → nmn-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nmn
-Version: 0.1.8
+Version: 0.1.10
 Summary: a neuron that matter
 Project-URL: Homepage, https://github.com/mlnomadpy/nmn
 Project-URL: Bug Tracker, https://github.com/mlnomadpy/my_package/issues

{nmn-0.1.8.dist-info → nmn-0.1.10.dist-info}/RECORD

@@ -3,7 +3,7 @@ nmn/keras/nmn.py,sha256=E7V7kyFB09PfMG1Da_TA2FirOiTCeAXYp3JWACV8h_c,5908
 nmn/linen/nmn.py,sha256=j4v6Z793wliE0xEAITde7jXu9Qras9u75NqdOSPSM4Q,3722
 nmn/nnx/TODO,sha256=U1WV51Eqij5igMjWLcbCjAZPONwIoPUQsMFKYHC6C8g,68
 nmn/nnx/nmn.py,sha256=tPNUtF8Lmv_B1TgMoVXfMQ9x0IPGKjSyAP6HnZ-YBsM,5651
-nmn/nnx/yatattention.py,sha256=4WyL9JW5wG05YURaLt76wA0zu2Bu2rMWtyFnVo9Gybo,24864
+nmn/nnx/yatattention.py,sha256=qEWiG_FIgr-TslYCbm2pcBi1myXJLC84nT6k1tMQcr4,25001
 nmn/nnx/yatconv.py,sha256=EOAAWfuv5QA-QTru-JyYKYNoGqxcklu7ph9a-CtmYsA,13123
 nmn/nnx/examples/language/mingpt.py,sha256=RveY3NwriTGPBdj8HNKDNtnXMaH0pgux8554m4Bhho4,61080
 nmn/nnx/examples/vision/cnn_cifar.py,sha256=UcK52-SCwuE2hl2BkpEbyg7N3Jwvvz8iFxiqhI7B9ew,73961
@@ -11,11 +11,11 @@ nmn/nnx/loss/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nmn/nnx/squashers/__init__.py,sha256=zXYPa3yzqMXxkIPvNHiaV6pcZRDOdVrzaVdYVDGALTY,180
 nmn/nnx/squashers/soft_tanh.py,sha256=WSJkxD6L9WU1eqPwsK2AW4V6OJbw5pSWYjKwkiWtLdo,812
 nmn/nnx/squashers/softer_sigmoid.py,sha256=vE6IWorZdBb2cww6fskARnwzdjTcWB2kKohuaJWVGNs,845
-nmn/nnx/squashers/softermax.py,sha256=NfxEDbogLUysyTvtVCTpDt27PplYvKRQLTZbYCL-Wfg,1226
+nmn/nnx/squashers/softermax.py,sha256=ggg0mHMFyk7b5xs31o-inNvWDzEvghD6YO3mtPlnkW4,1318
 nmn/tf/nmn.py,sha256=A-K65z9_aN62tAy12b0553nXxrzOofK1umGMRGJYjqw,6036
 nmn/torch/conv.py,sha256=g5YxStk1p85WkvfecqbzRZaWaAJahOSArpMcqxWAWKc,83413
 nmn/torch/nmn.py,sha256=8K0S3nwpGprT7apbCqpaYpKpxq8F8g8EL8PHIezgMCY,4658
-nmn-0.1.8.dist-info/METADATA,sha256=8DvDHl3Tkp3HSge5rYHZdKaKTso3j-H1lDgF3c4owrI,8800
-nmn-0.1.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-nmn-0.1.8.dist-info/licenses/LICENSE,sha256=kbZSd5WewnN2PSjvAC6DprP7pXx6NUNsnltmU2Mz1yA,34519
-nmn-0.1.8.dist-info/RECORD,,
+nmn-0.1.10.dist-info/METADATA,sha256=o-wLjeO-n2h56-cvw-AqrRiio5UFaerm58w03XkdHQY,8801
+nmn-0.1.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+nmn-0.1.10.dist-info/licenses/LICENSE,sha256=kbZSd5WewnN2PSjvAC6DprP7pXx6NUNsnltmU2Mz1yA,34519
+nmn-0.1.10.dist-info/RECORD,,