nmn 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

nmn/nnx/squashers/__init__.py

@@ -0,0 +1,9 @@
+from .softermax import softermax
+from .softer_sigmoid import softer_sigmoid
+from .soft_tanh import soft_tanh
+
+__all__ = [
+    "softermax",
+    "softer_sigmoid",
+    "soft_tanh",
+]

nmn/nnx/squashers/soft_tanh.py

@@ -0,0 +1,29 @@
+import jax.numpy as jnp
+from jax import Array
+
+def soft_tanh(
+    x: Array,
+    n: float = 1.0,
+) -> Array:
+    """
+    Maps a non-negative score to the range [-1, 1) using the soft-tanh function.
+
+    The soft-tanh function is defined as:
+    .. math::
+        \\text{soft-tanh}_n(x) = \\frac{x^n - 1}{1 + x^n}
+
+    The power `n` again controls the transition sharpness: higher `n` makes the
+    function approach -1 more quickly for large `x`.
+
+    Args:
+        x (Array): A JAX array of non-negative scores (x >= 0).
+        n (float, optional): The power to raise the score to. Defaults to 1.0.
+
+    Returns:
+        Array: The mapped scores in the range [-1, 1).
+    """
+    if n <= 0:
+        raise ValueError("Power 'n' must be positive.")
+
+    x_n = jnp.power(x, n)
+    return (x_n - 1.0) / (1.0 + x_n)

nmn/nnx/squashers/softer_sigmoid.py

@@ -0,0 +1,29 @@
+import jax.numpy as jnp
+from jax import Array
+
+def softer_sigmoid(
+    x: Array,
+    n: float = 1.0,
+) -> Array:
+    """
+    Squashes a non-negative score into the range [0, 1) using the soft-sigmoid function.
+
+    The soft-sigmoid function is defined as:
+    .. math::
+        \\text{soft-sigmoid}_n(x) = \\frac{x^n}{1 + x^n}
+    
+    The power `n` modulates the softness: higher `n` makes the function approach
+    zero faster for large `x`, while `n < 1` makes the decay slower.
+
+    Args:
+        x (Array): A JAX array of non-negative scores (x >= 0).
+        n (float, optional): The power to raise the score to. Defaults to 1.0.
+
+    Returns:
+        Array: The squashed scores in the range [0, 1).
+    """
+    if n <= 0:
+        raise ValueError("Power 'n' must be positive.")
+    
+    x_n = jnp.power(x, n)
+    return x_n / (1.0 + x_n)

nmn/nnx/squashers/softermax.py

@@ -0,0 +1,38 @@
+import jax.numpy as jnp
+from jax import Array
+from typing import Optional
+
+def softermax(
+    x: Array,
+    n: float = 1.0,
+    epsilon: float = 1e-12,
+    axis: Optional[int] = -1,
+) -> Array:
+    """
+    Normalizes a set of non-negative scores using the Softermax function.
+
+    The Softermax function is defined as:
+    .. math::
+        \\text{softermax}_n(x_k, \\{x_i\\}) = \\frac{x_k^n}{\\epsilon + \\sum_i x_i^n}
+
+    The power `n` controls the sharpness of the distribution: `n=1` recovers
+    the original Softermax, while `n > 1` makes the distribution harder (more
+    peaked), and `0 < n < 1` makes it softer.
+
+    Args:
+        x (Array): A JAX array of non-negative scores.
+        n (float, optional): The power to raise each score to. Defaults to 1.0.
+        epsilon (float, optional): A small constant for numerical stability.
+            Defaults to 1e-12.
+        axis (Optional[int], optional): The axis to perform the sum over.
+            Defaults to -1.
+
+    Returns:
+        Array: The normalized scores.
+    """
+    if n <= 0:
+        raise ValueError("Power 'n' must be positive.")
+
+    x_n = jnp.power(x, n)
+    sum_x_n = jnp.sum(x_n, axis=axis, keepdims=True)
+    return x_n / (epsilon + sum_x_n)

nmn/nnx/yatattention.py

@@ -26,7 +26,8 @@ from flax.typing import (
   DotGeneralT,
 )
 
-
+from nmn.nnx.nmn import YatNMN
+from jax import Array
 
 def yat_attention_weights(
   query: Array,
@@ -153,91 +154,6 @@ def yat_attention(
     '...hqk,...khd->...qhd', attn_weights, value, precision=precision
   )
 
-Array = jax.Array
-
-# Add YatNMN class implementation
-default_bias_init = initializers.zeros_init()
-default_alpha_init = initializers.ones_init()
-
-class YatNMN(Module):
-  """A linear transformation with custom distance-based computation."""
-
-  def __init__(
-    self,
-    in_features: int,
-    out_features: int,
-    *,
-    use_bias: bool = True,
-    use_alpha: bool = True,
-    dtype: Optional[Dtype] = None,
-    param_dtype: Dtype = jnp.float32,
-    precision: PrecisionLike = None,
-    kernel_init: Initializer = default_kernel_init,
-    bias_init: Initializer = default_bias_init,
-    alpha_init: Initializer = default_alpha_init,
-    dot_general: DotGeneralT = lax.dot_general,
-    rngs: rnglib.Rngs,
-    epsilon: float = 1e-5,
-  ):
-
-    kernel_key = rngs.params()
-    self.kernel = nnx.Param(
-      kernel_init(kernel_key, (in_features, out_features), param_dtype)
-    )
-    self.bias: nnx.Param[jax.Array] | None
-    if use_bias:
-      bias_key = rngs.params()
-      self.bias = nnx.Param(bias_init(bias_key, (out_features,), param_dtype))
-    else:
-      self.bias = None
-
-    self.alpha: nnx.Param[jax.Array] | None
-    if use_alpha:
-      alpha_key = rngs.params()
-      self.alpha = nnx.Param(alpha_init(alpha_key, (1,), param_dtype))
-    else:
-      self.alpha = None
-
-    self.in_features = in_features
-    self.out_features = out_features
-    self.use_bias = use_bias
-    self.use_alpha = use_alpha
-    self.dtype = dtype
-    self.param_dtype = param_dtype
-    self.precision = precision
-    self.kernel_init = kernel_init
-    self.bias_init = bias_init
-    self.dot_general = dot_general
-    self.epsilon = epsilon
-
-  def __call__(self, inputs: Array) -> Array:
-    """Applies YatNMN transformation to inputs."""
-    kernel = self.kernel.value
-    bias = self.bias.value if self.bias is not None else None
-    alpha = self.alpha.value if self.alpha is not None else None
-
-    y = self.dot_general(
-      inputs,
-      kernel,
-      (((inputs.ndim - 1,), (0,)), ((), ())),
-      precision=self.precision,
-    )
-
-    inputs_squared_sum = jnp.sum(inputs**2, axis=-1, keepdims=True)
-    kernel_squared_sum = jnp.sum(kernel**2, axis=0, keepdims=True)
-    distances = inputs_squared_sum + kernel_squared_sum - 2 * y
-
-    # Element-wise operation
-    y = y ** 2 / (distances + self.epsilon)
-
-    if bias is not None:
-      y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))
-
-    if alpha is not None:
-      scale = (jnp.sqrt(self.out_features) / jnp.log(1 + self.out_features)) ** alpha
-      y = y * scale
-
-    return y
 
 
 def dot_product_attention_weights(
@@ -435,6 +351,10 @@ class MultiHeadAttention(Module):
     attention_fn: Callable[..., Array] = yat_attention,
     decode: bool | None = None,
     normalize_qk: bool = False,
+    use_alpha: bool = True,
+    alpha_init: Initializer = initializers.ones_init(),
+    use_dropconnect: bool = False,
+    dropconnect_rate: float = 0.0,
     # Deprecated, will be removed.
     qkv_dot_general: DotGeneralT | None = None,
     out_dot_general: DotGeneralT | None = None,
@@ -470,6 +390,10 @@ class MultiHeadAttention(Module):
     self.qkv_dot_general_cls = qkv_dot_general_cls
     self.out_dot_general_cls = out_dot_general_cls
     self.epsilon = epsilon
+    self.use_alpha = use_alpha
+    self.alpha_init = alpha_init
+    self.use_dropconnect = use_dropconnect
+    self.dropconnect_rate = dropconnect_rate
 
     if self.qkv_features % self.num_heads != 0:
       raise ValueError(
@@ -491,6 +415,10 @@ class MultiHeadAttention(Module):
       use_bias=self.use_bias,
       precision=self.precision,
       epsilon=self.epsilon,
+      use_alpha=self.use_alpha,
+      alpha_init=self.alpha_init,
+      use_dropconnect=self.use_dropconnect,
+      drop_rate=self.dropconnect_rate,
     )
 
     # project inputs_q to multi-headed q/k/v
@@ -590,10 +518,23 @@ class MultiHeadAttention(Module):
         f'but module expects {self.in_features}.'
       )
 
+    is_deterministic: bool = False
+    if self.dropout_rate > 0.0 or (
+      self.use_dropconnect and self.dropconnect_rate > 0.0
+    ):
+      is_deterministic = first_from(
+        deterministic,
+        self.deterministic,
+        error_msg="""No `deterministic` argument was provided to MultiHeadAttention
+          as either a __call__ argument, class attribute, or nnx.flag.""",
+      )
+    else:
+      is_deterministic = True
+
     # Apply YatNMN transformations and reshape to multi-head format
-    query = squash(self.query(inputs_q))
-    key = squash(self.key(inputs_k))
-    value = squash(self.value(inputs_v))
+    query = self.query(inputs_q, deterministic=is_deterministic)
+    key = self.key(inputs_k, deterministic=is_deterministic)
+    value = self.value(inputs_v, deterministic=is_deterministic)
 
     # Reshape from [batch..., length, qkv_features] to [batch..., length, num_heads, head_dim]
     query = query.reshape(query.shape[:-1] + (self.num_heads, self.head_dim))
@@ -660,26 +601,11 @@ class MultiHeadAttention(Module):
         ),
       )
 
-    if (
-      self.dropout_rate > 0.0
-    ):  # Require `deterministic` only if using dropout.
-      deterministic = first_from(
-        deterministic,
-        self.deterministic,
-        error_msg="""No `deterministic` argument was provided to MultiHeadAttention
-          as either a __call__ argument, class attribute, or nnx.flag.""",
-      )
-      if not deterministic:
-        if rngs is None:
-          raise ValueError(
-            "'rngs' must be provided if 'dropout_rng' is not given."
-          )
-        dropout_rng = rngs.dropout()
-      else:
-        dropout_rng = None
-    else:
-      deterministic = True
-      dropout_rng = None
+    dropout_rng = None
+    if self.dropout_rate > 0.0 and not is_deterministic:
+      if rngs is None:
+        raise ValueError("'rngs' must be provided for dropout.")
+      dropout_rng = rngs.dropout()
 
     # apply attention with epsilon parameter for YatNMN
     x = self.attention_fn(
@@ -690,7 +616,7 @@ class MultiHeadAttention(Module):
       dropout_rng=dropout_rng,
       dropout_rate=self.dropout_rate,
       broadcast_dropout=self.broadcast_dropout,
-      deterministic=deterministic,
+      deterministic=is_deterministic,
       dtype=self.dtype,
       precision=self.precision,
       module=self if sow_weights else None,

{nmn-0.1.5.dist-info → nmn-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nmn
-Version: 0.1.5
+Version: 0.1.6
 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.6.dist-info/RECORD

@@ -0,0 +1,19 @@
+nmn/__init__.py,sha256=F_5o-lCggdEdWfR1l1YC_jfR01mJmveugwUndoRx8n8,83
+nmn/keras/nmn.py,sha256=E7V7kyFB09PfMG1Da_TA2FirOiTCeAXYp3JWACV8h_c,5908
+nmn/linen/nmn.py,sha256=j4v6Z793wliE0xEAITde7jXu9Qras9u75NqdOSPSM4Q,3722
+nmn/nnx/nmn.py,sha256=tPNUtF8Lmv_B1TgMoVXfMQ9x0IPGKjSyAP6HnZ-YBsM,5651
+nmn/nnx/yatattention.py,sha256=i6XfCGHISyb2P6KrgYFnhhdzqSTWAyshFhy1XEeuEWc,24642
+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
+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/tf/nmn.py,sha256=A-K65z9_aN62tAy12b0553nXxrzOofK1umGMRGJYjqw,6036
+nmn/torch/nmn.py,sha256=8K0S3nwpGprT7apbCqpaYpKpxq8F8g8EL8PHIezgMCY,4658
+nmn-0.1.6.dist-info/METADATA,sha256=Y9MByC16wz1MGYVZRmZA0wJQATB6Kj6w6TOL5lPzl0Q,8800
+nmn-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+nmn-0.1.6.dist-info/licenses/LICENSE,sha256=kbZSd5WewnN2PSjvAC6DprP7pXx6NUNsnltmU2Mz1yA,34519
+nmn-0.1.6.dist-info/RECORD,,

nmn-0.1.5.dist-info/RECORD

@@ -1,14 +0,0 @@
-nmn/__init__.py,sha256=F_5o-lCggdEdWfR1l1YC_jfR01mJmveugwUndoRx8n8,83
-nmn/keras/nmn.py,sha256=E7V7kyFB09PfMG1Da_TA2FirOiTCeAXYp3JWACV8h_c,5908
-nmn/linen/nmn.py,sha256=j4v6Z793wliE0xEAITde7jXu9Qras9u75NqdOSPSM4Q,3722
-nmn/nnx/nmn.py,sha256=tPNUtF8Lmv_B1TgMoVXfMQ9x0IPGKjSyAP6HnZ-YBsM,5651
-nmn/nnx/yatattention.py,sha256=chjtUKJtaR7ROPnNqkicbvMs7hzZKE0fIo_8cTNiju8,26601
-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
-nmn/tf/nmn.py,sha256=A-K65z9_aN62tAy12b0553nXxrzOofK1umGMRGJYjqw,6036
-nmn/torch/nmn.py,sha256=8K0S3nwpGprT7apbCqpaYpKpxq8F8g8EL8PHIezgMCY,4658
-nmn-0.1.5.dist-info/METADATA,sha256=7gvXle6Hgdgyj_tJk1DGdkOh03BOsfSks-ZHPOIEwHQ,8800
-nmn-0.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-nmn-0.1.5.dist-info/licenses/LICENSE,sha256=kbZSd5WewnN2PSjvAC6DprP7pXx6NUNsnltmU2Mz1yA,34519
-nmn-0.1.5.dist-info/RECORD,,