nmn 0.1.0__py3-none-any.whl

nmn/__init__.py

@@ -0,0 +1,3 @@
+"""Neural-Matter Network (NMN) - beyond blinded neurons."""
+
+__version__ = "0.1.0"

nmn/keras/nmn.py

@@ -0,0 +1,153 @@
+from keras.src import activations, constraints, initializers, regularizers
+from keras.src.api_export import keras_export
+from keras.src.layers.input_spec import InputSpec
+from keras.src.layers.layer import Layer
+from keras.src import ops
+import math
+
+@keras_export("keras.layers.YatDense")
+class YatNMN(Layer):
+    """A YAT densely-connected NN layer.
+
+    This layer implements the operation:
+    `output = scale * (dot(input, kernel)^2 / (squared_euclidean_distance + epsilon))`
+    where:
+    - `scale` is a dynamic scaling factor based on output dimension
+    - `squared_euclidean_distance` is computed between input and kernel
+    - `epsilon` is a small constant to prevent division by zero
+
+    Args:
+        units: Positive integer, dimensionality of the output space.
+        use_bias: Boolean, whether the layer uses a bias vector.
+        epsilon: Float, small constant added to denominator for numerical stability.
+        kernel_initializer: Initializer for the `kernel` weights matrix.
+        bias_initializer: Initializer for the bias vector.
+        kernel_regularizer: Regularizer function applied to the `kernel` weights matrix.
+        bias_regularizer: Regularizer function applied to the bias vector.
+        activity_regularizer: Regularizer function applied to the output.
+        kernel_constraint: Constraint function applied to the `kernel` weights matrix.
+        bias_constraint: Constraint function applied to the bias vector.
+
+    Input shape:
+        N-D tensor with shape: `(batch_size, ..., input_dim)`.
+        The most common situation would be a 2D input with shape
+        `(batch_size, input_dim)`.
+
+    Output shape:
+        N-D tensor with shape: `(batch_size, ..., units)`.
+        For instance, for a 2D input with shape `(batch_size, input_dim)`,
+        the output would have shape `(batch_size, units)`.
+    """
+
+    def __init__(
+        self,
+        units,
+        activation=None,
+        use_bias=True,
+        epsilon=1e-5,
+        kernel_initializer="orthogonal",
+        bias_initializer="zeros",
+        kernel_regularizer=None,
+        bias_regularizer=None,
+        activity_regularizer=None,
+        kernel_constraint=None,
+        bias_constraint=None,
+        **kwargs,
+    ):
+        super().__init__(activity_regularizer=activity_regularizer, **kwargs)
+        self.units = units
+        self.activation = activations.get(activation)
+        self.use_bias = use_bias
+        self.epsilon = epsilon
+
+        self.kernel_initializer = initializers.get(kernel_initializer)
+        self.bias_initializer = initializers.get(bias_initializer)
+        self.kernel_regularizer = regularizers.get(kernel_regularizer)
+        self.bias_regularizer = regularizers.get(bias_regularizer)
+        self.kernel_constraint = constraints.get(kernel_constraint)
+        self.bias_constraint = constraints.get(bias_constraint)
+
+        self.input_spec = InputSpec(min_ndim=2)
+        self.supports_masking = True
+
+    def build(self, input_shape):
+        input_dim = input_shape[-1]
+
+        self.kernel = self.add_weight(
+            name="kernel",
+            shape=(input_dim, self.units),
+            initializer=self.kernel_initializer,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+            trainable=True,
+        )
+
+        if self.use_bias:
+            self.bias = self.add_weight(
+                name="bias",
+                shape=(self.units,),
+                initializer=self.bias_initializer,
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+                trainable=True,
+            )
+        else:
+            self.bias = None
+
+        # Add alpha parameter for dynamic scaling
+        self.alpha = self.add_weight(
+            name="alpha",
+            shape=(1,),
+            initializer="ones",
+            trainable=True,
+        )
+
+        self.input_spec = InputSpec(min_ndim=2, axes={-1: input_dim})
+        self.built = True
+
+    def call(self, inputs):
+        # Compute dot product
+        dot_product = ops.matmul(inputs, self.kernel)
+
+        # Compute squared distances
+        inputs_squared_sum = ops.sum(inputs ** 2, axis=-1, keepdims=True)
+        kernel_squared_sum = ops.sum(self.kernel ** 2, axis=0)
+        distances = inputs_squared_sum + kernel_squared_sum - 2 * dot_product
+
+        # Compute inverse square attention
+        outputs = dot_product ** 2 / (distances + self.epsilon)
+        if self.use_bias:
+            outputs = ops.add(outputs, self.bias)
+
+        # Apply dynamic scaling
+        scale = (ops.sqrt(ops.cast(self.units, self.compute_dtype)) /
+                ops.log1p(ops.cast(self.units, self.compute_dtype))) ** self.alpha
+        outputs = outputs * scale
+
+
+        if self.activation is not None:
+            outputs = self.activation(outputs)
+
+        return outputs
+
+    def compute_output_shape(self, input_shape):
+        output_shape = list(input_shape)
+        output_shape[-1] = self.units
+        return tuple(output_shape)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            "units": self.units,
+            "activation": activations.serialize(self.activation),
+            "use_bias": self.use_bias,
+            "epsilon": self.epsilon,
+            "kernel_initializer": initializers.serialize(self.kernel_initializer),
+            "bias_initializer": initializers.serialize(self.bias_initializer),
+            "kernel_regularizer": regularizers.serialize(self.kernel_regularizer),
+            "bias_regularizer": regularizers.serialize(self.bias_regularizer),
+            "activity_regularizer": regularizers.serialize(self.activity_regularizer),
+            "kernel_constraint": constraints.serialize(self.kernel_constraint),
+            "bias_constraint": constraints.serialize(self.bias_constraint),
+        })
+        return config

nmn/linen/nmn.py

@@ -0,0 +1,112 @@
+import jax.numpy as jnp
+from flax.linen.dtypes import promote_dtype
+from flax.linen.module import Module, compact
+from flax.typing import (
+  PRNGKey as PRNGKey,
+  Shape as Shape,
+  DotGeneralT,
+)
+
+from typing import (
+  Any,
+)
+import jax.numpy as jnp
+import jax.lax as lax
+from flax.linen import Module, compact
+from flax import linen as nn
+from flax.linen.initializers import zeros_init, lecun_normal
+from typing import Any, Optional
+
+class YatNMN(Module):
+    """A custom transformation applied over the last dimension of the input using squared Euclidean distance.
+
+    Attributes:
+      features: the number of output features.
+      use_bias: whether to add a bias to the output (default: True).
+      dtype: the dtype of the computation (default: infer from input and params).
+      param_dtype: the dtype passed to parameter initializers (default: float32).
+      precision: numerical precision of the computation see ``jax.lax.Precision`` for details.
+      kernel_init: initializer function for the weight matrix.
+      bias_init: initializer function for the bias.
+      epsilon: small constant added to avoid division by zero (default: 1e-6).
+    """
+    features: int
+    use_bias: bool = True
+    use_alpha: bool = True
+    dtype: Optional[Any] = None
+    param_dtype: Any = jnp.float32
+    precision: Any = None
+    kernel_init: Any = nn.initializers.orthogonal()
+    bias_init: Any = zeros_init()
+
+    alpha_init: Any = lambda key, shape, dtype: jnp.ones(shape, dtype)  # Initialize alpha to 1.0
+    epsilon: float = 1e-6
+    dot_general: DotGeneralT | None = None
+    dot_general_cls: Any = None
+    return_weights: bool = False
+
+    @compact
+    def __call__(self, inputs: Any) -> Any:
+        """Applies a transformation to the inputs along the last dimension using squared Euclidean distance.
+
+        Args:
+          inputs: The nd-array to be transformed.
+
+        Returns:
+          The transformed input.
+        """
+        kernel = self.param(
+            'kernel',
+            self.kernel_init,
+            (self.features, jnp.shape(inputs)[-1]),
+            self.param_dtype,
+        )
+        if self.use_alpha:
+            alpha = self.param(
+                'alpha',
+                self.alpha_init,
+                (1,),  # Single scalar parameter
+                self.param_dtype,
+            )
+        else:
+            alpha = None
+
+        if self.use_bias:
+            bias = self.param(
+                'bias', self.bias_init, (self.features,), self.param_dtype
+            )
+        else:
+            bias = None
+
+        inputs, kernel, bias, alpha = promote_dtype(inputs, kernel, bias, alpha, dtype=self.dtype)
+
+        # Compute dot product between input and kernel
+        if self.dot_general_cls is not None:
+          dot_general = self.dot_general_cls()
+        elif self.dot_general is not None:
+          dot_general = self.dot_general
+        else:
+          dot_general = lax.dot_general
+        y = dot_general(
+          inputs,
+          jnp.transpose(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=-1)
+        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.features) / jnp.log(1 + self.features)) ** alpha
+          y = y * scale
+
+        # Normalize y
+        if self.return_weights:
+           return y, kernel
+        return y

nmn/nnx/nmn.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import typing as tp
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import lax
+import opt_einsum
+
+from flax.core.frozen_dict import FrozenDict
+from flax import nnx
+from flax.nnx import rnglib, variablelib
+from flax.nnx.module import Module, first_from
+from flax.nnx.nn import dtypes, initializers
+from flax.typing import (
+  Dtype,
+  Shape,
+  Initializer,
+  PrecisionLike,
+  DotGeneralT,
+  ConvGeneralDilatedT,
+  PaddingLike,
+  LaxPadding,
+  PromoteDtypeFn,
+  EinsumT,
+)
+
+Array = jax.Array
+Axis = int
+Size = int
+
+
+default_kernel_init = initializers.lecun_normal()
+default_bias_init = initializers.zeros_init()
+default_alpha_init = initializers.ones_init()
+
+class YatNMN(Module):
+  """A linear transformation applied over the last dimension of the input.
+
+  Example usage::
+
+    >>> from flax import nnx
+    >>> import jax, jax.numpy as jnp
+
+    >>> layer = nnx.Linear(in_features=3, out_features=4, rngs=nnx.Rngs(0))
+    >>> jax.tree.map(jnp.shape, nnx.state(layer))
+    State({
+      'bias': VariableState(
+        type=Param,
+        value=(4,)
+      ),
+      'kernel': VariableState(
+        type=Param,
+        value=(3, 4)
+      )
+    })
+
+  Args:
+    in_features: the number of input features.
+    out_features: the number of output features.
+    use_bias: whether to add a bias to the output (default: True).
+    dtype: the dtype of the computation (default: infer from input and params).
+    param_dtype: the dtype passed to parameter initializers (default: float32).
+    precision: numerical precision of the computation see ``jax.lax.Precision``
+      for details.
+    kernel_init: initializer function for the weight matrix.
+    bias_init: initializer function for the bias.
+    dot_general: dot product function.
+    promote_dtype: function to promote the dtype of the arrays to the desired
+      dtype. The function should accept a tuple of ``(inputs, kernel, bias)``
+      and a ``dtype`` keyword argument, and return a tuple of arrays with the
+      promoted dtype.
+    rngs: rng key.
+  """
+
+  __data__ = ('kernel', 'bias')
+
+  def __init__(
+    self,
+    in_features: int,
+    out_features: int,
+    *,
+    use_bias: bool = True,
+    use_alpha: bool = True,
+    dtype: tp.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,
+    promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,
+    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(bias_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.promote_dtype = promote_dtype
+    self.epsilon = epsilon
+
+  def __call__(self, inputs: Array) -> Array:
+    """Applies a linear transformation to the inputs along the last dimension.
+
+    Args:
+      inputs: The nd-array to be transformed.
+
+    Returns:
+      The transformed input.
+    """
+    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
+
+    inputs, kernel, bias, alpha = self.promote_dtype(
+      (inputs, kernel, bias, alpha), dtype=self.dtype
+    )
+    y = self.dot_general(
+      inputs,
+      kernel,
+      (((inputs.ndim - 1,), (0,)), ((), ())),
+      precision=self.precision,
+    )
+
+    assert self.use_bias == (bias is not None)
+    assert self.use_alpha == (alpha is not None)
+
+    inputs_squared_sum = jnp.sum(inputs**2, axis=-1, keepdims=True)
+    kernel_squared_sum = jnp.sum(kernel**2, axis=0, keepdims=True)  # Change axis to 0 and keepdims to 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