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

nmn/nnx/nmn.py

@@ -4,26 +4,18 @@ 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 import rnglib
+from flax.nnx.module import Module
 from flax.nnx.nn import dtypes, initializers
 from flax.typing import (
   Dtype,
-  Shape,
   Initializer,
   PrecisionLike,
   DotGeneralT,
-  ConvGeneralDilatedT,
-  PaddingLike,
-  LaxPadding,
   PromoteDtypeFn,
-  EinsumT,
 )
 
 Array = jax.Array
@@ -60,21 +52,26 @@ class YatNMN(Module):
     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).
+    use_alpha: whether to use alpha scaling (default: True).
+    use_dropconnect: whether to use DropConnect (default: False).
     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.
+    alpha_init: initializer function for the alpha.
     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.
+    epsilon: A small float added to the denominator to prevent division by zero.
+    drop_rate: dropout rate for DropConnect (default: 0.0).
     rngs: rng key.
   """
 
-  __data__ = ('kernel', 'bias')
+  __data__ = ('kernel', 'bias', 'alpha', 'dropconnect_key')
 
   def __init__(
     self,
@@ -83,6 +80,7 @@ class YatNMN(Module):
     *,
     use_bias: bool = True,
     use_alpha: bool = True,
+    use_dropconnect: bool = False,
     dtype: tp.Optional[Dtype] = None,
     param_dtype: Dtype = jnp.float32,
     precision: PrecisionLike = None,
@@ -91,8 +89,9 @@ class YatNMN(Module):
     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,
+    drop_rate: float = 0.0,
+    rngs: rnglib.Rngs,
   ):
 
     kernel_key = rngs.params()
@@ -117,6 +116,7 @@ class YatNMN(Module):
     self.out_features = out_features
     self.use_bias = use_bias
     self.use_alpha = use_alpha
+    self.use_dropconnect = use_dropconnect
     self.dtype = dtype
     self.param_dtype = param_dtype
     self.precision = precision
@@ -125,12 +125,19 @@ class YatNMN(Module):
     self.dot_general = dot_general
     self.promote_dtype = promote_dtype
     self.epsilon = epsilon
+    self.drop_rate = drop_rate
+
+    if use_dropconnect:
+      self.dropconnect_key = rngs.params()
+    else:
+      self.dropconnect_key = None
 
-  def __call__(self, inputs: Array) -> Array:
+  def __call__(self, inputs: Array, *, deterministic: bool = False) -> Array:
     """Applies a linear transformation to the inputs along the last dimension.
 
     Args:
       inputs: The nd-array to be transformed.
+      deterministic: If true, DropConnect is not applied (e.g., during inference).
 
     Returns:
       The transformed input.
@@ -139,6 +146,11 @@ class YatNMN(Module):
     bias = self.bias.value if self.bias is not None else None
     alpha = self.alpha.value if self.alpha is not None else None
 
+    if self.use_dropconnect and not deterministic and self.drop_rate > 0.0:
+      keep_prob = 1.0 - self.drop_rate
+      mask = jax.random.bernoulli(self.dropconnect_key, p=keep_prob, shape=kernel.shape)
+      kernel = (kernel * mask) / keep_prob
+
     inputs, kernel, bias, alpha = self.promote_dtype(
       (inputs, kernel, bias, alpha), dtype=self.dtype
     )
@@ -166,5 +178,4 @@ class YatNMN(Module):
       scale = (jnp.sqrt(self.out_features) / jnp.log(1 + self.out_features)) ** alpha
       y = y * scale
 
-
     return y

nmn/nnx/yatconv.py

@@ -110,6 +110,8 @@ class YatConv(Module):
     feature_group_count: integer, default 1. If specified divides the input
       features into groups.
     use_bias: whether to add a bias to the output (default: True).
+    use_alpha: whether to use alpha scaling (default: True).
+    use_dropconnect: whether to use DropConnect (default: False).
     mask: Optional mask for the weights during masked convolution. The mask must
           be the same shape as the convolution weight matrix.
     dtype: the dtype of the computation (default: infer from input and params).
@@ -123,10 +125,11 @@ class YatConv(Module):
       and a ``dtype`` keyword argument, and return a tuple of arrays with the
       promoted dtype.
     epsilon: A small float added to the denominator to prevent division by zero.
+    drop_rate: dropout rate for DropConnect (default: 0.0).
     rngs: rng key.
   """
 
-  __data__ = ('kernel', 'bias', 'mask')
+  __data__ = ('kernel', 'bias', 'mask', 'dropconnect_key')
 
   def __init__(
     self,
@@ -142,6 +145,7 @@ class YatConv(Module):
     
     use_bias: bool = True,
     use_alpha: bool = True,
+    use_dropconnect: bool = False,
     kernel_init: Initializer = default_kernel_init,
     bias_init: Initializer = default_bias_init,
     alpha_init: Initializer = default_alpha_init,
@@ -153,6 +157,7 @@ class YatConv(Module):
     conv_general_dilated: ConvGeneralDilatedT = lax.conv_general_dilated,
     promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,
     epsilon: float = 1e-5,
+    drop_rate: float = 0.0,
     rngs: rnglib.Rngs,
   ):
     if isinstance(kernel_size, int):
@@ -185,6 +190,7 @@ class YatConv(Module):
     self.feature_group_count = feature_group_count
     self.use_bias = use_bias
     self.use_alpha = use_alpha
+    self.use_dropconnect = use_dropconnect
 
     self.mask = mask 
     self.dtype = dtype
@@ -195,6 +201,7 @@ class YatConv(Module):
     self.conv_general_dilated = conv_general_dilated
     self.promote_dtype = promote_dtype
     self.epsilon = epsilon
+    self.drop_rate = drop_rate
 
     if use_alpha:
       alpha_key = rngs.params()
@@ -202,8 +209,12 @@ class YatConv(Module):
     else:
       self.alpha = None
 
+    if use_dropconnect:
+      self.dropconnect_key = rngs.params()
+    else:
+      self.dropconnect_key = None
 
-  def __call__(self, inputs: Array) -> Array:
+  def __call__(self, inputs: Array, *, deterministic: bool = False) -> Array:
     assert isinstance(self.kernel_size, tuple)
 
     def maybe_broadcast(
@@ -261,6 +272,12 @@ class YatConv(Module):
 
     kernel_val = self.kernel.value
     
+    # Apply DropConnect if enabled and not in deterministic mode
+    if self.use_dropconnect and not deterministic and self.drop_rate > 0.0:
+      keep_prob = 1.0 - self.drop_rate
+      mask = jax.random.bernoulli(self.dropconnect_key, p=keep_prob, shape=kernel_val.shape)
+      kernel_val = (kernel_val * mask) / keep_prob
+    
     current_mask = self.mask 
     if current_mask is not None:
       if current_mask.shape != self.kernel_shape:

nmn-0.1.5.dist-info/METADATA

@@ -0,0 +1,176 @@
+Metadata-Version: 2.4
+Name: nmn
+Version: 0.1.5
+Summary: a neuron that matter
+Project-URL: Homepage, https://github.com/mlnomadpy/nmn
+Project-URL: Bug Tracker, https://github.com/mlnomadpy/my_package/issues
+Author-email: Taha Bouhsine <yat@mlnomads.com>
+License-File: LICENSE
+Classifier: License :: OSI Approved :: GNU Affero General Public License v3
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+
+# nmn
+Not the neurons we want, but the neurons we need
+
+[![PyPI version](https://img.shields.io/pypi/v/nmn.svg)](https://pypi.org/project/nmn/)
+[![Downloads](https://static.pepy.tech/badge/nmn)](https://pepy.tech/project/nmn)
+[![Downloads/month](https://static.pepy.tech/badge/nmn/month)](https://pepy.tech/project/nmn)
+[![GitHub stars](https://img.shields.io/github/stars/mlnomadpy/nmn?style=social)](https://github.com/mlnomadpy/nmn)
+[![GitHub forks](https://img.shields.io/github/forks/mlnomadpy/nmn?style=social)](https://github.com/mlnomadpy/nmn)
+[![GitHub issues](https://img.shields.io/github/issues/mlnomadpy/nmn)](https://github.com/mlnomadpy/nmn/issues)
+[![PyPI - License](https://img.shields.io/pypi/l/nmn)](https://pypi.org/project/nmn/)
+[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/nmn)](https://pypi.org/project/nmn/)
+
+## Features
+
+*   **Activation-Free Non-linearity:** Learns complex, non-linear relationships without separate activation functions.
+*   **Multiple Frameworks:** Supports Flax (Linen & NNX), Keras, PyTorch, and TensorFlow.
+*   **Yat-Product & Yat-Conv:** Implements novel Yat-Product and Yat-Conv operations.
+*   **Inspired by Research:** Based on the principles from "Deep Learning 2.0/2.1: Artificial Neurons that Matter".
+
+## Overview
+
+**nmn** provides neural network layers for multiple frameworks (Flax, NNX, Keras, PyTorch, TensorFlow) that do not require activation functions to learn non-linearity. The main goal is to enable deep learning architectures where the layer itself is inherently non-linear, inspired by the papers:
+
+> Deep Learning 2.0: Artificial Neurons that Matter: Reject Correlation - Embrace Orthogonality
+>
+> Deep Learning 2.1: Deep Learning 2.1: Mind and Cosmos - Towards Cosmos-Inspired Interpretable Neural Networks
+
+## Math
+
+Yat-Product: 
+$$
+ⵟ(\mathbf{w},\mathbf{x}) := \frac{\langle \mathbf{w}, \mathbf{x} \rangle^2}{\|\mathbf{w} - \mathbf{x}\|^2 + \epsilon} = \frac{ \|\mathbf{x}\|^2  \|\mathbf{w}\|^2 \cos^2 \theta}{\|\mathbf{w}\|^2 - 2\mathbf{w}^\top\mathbf{x} + \|\mathbf{x}\|^2 + \epsilon} = \frac{ \|\mathbf{x}\|^2  \|\mathbf{w}\|^2 \cos^2 \theta}{((\mathbf{x}-\mathbf{w})\cdot(\mathbf{x}-\mathbf{w}))^2 + \epsilon}.
+$$
+
+**Explanation:**
+- $\mathbf{w}$ is the weight vector, $\mathbf{x}$ is the input vector.
+- $\langle \mathbf{w}, \mathbf{x} \rangle$ is the dot product between $\mathbf{w}$ and $\mathbf{x}$.
+- $\|\mathbf{w} - \mathbf{x}\|^2$ is the squared Euclidean distance between $\mathbf{w}$ and $\mathbf{x}$.
+- $\epsilon$ is a small constant for numerical stability.
+- $\theta$ is the angle between $\mathbf{w}$ and $\mathbf{x}$.
+
+This operation:
+- **Numerator:** Squares the similarity (dot product) between $\mathbf{w}$ and $\mathbf{x}$, emphasizing strong alignments.
+- **Denominator:** Penalizes large distances, so the response is high only when $\mathbf{w}$ and $\mathbf{x}$ are both similar in direction and close in space.
+- **No activation needed:** The non-linearity is built into the operation itself, allowing the layer to learn complex, non-linear relationships without a separate activation function.
+- **Geometric view:** The output is maximized when $\mathbf{w}$ and $\mathbf{x}$ are both large in norm, closely aligned (small $\theta$), and close together in Euclidean space.
+
+Yat-Conv:
+$$
+ⵟ^*(\mathbf{W}, \mathbf{X}) := \frac{\langle \mathbf{w}, \mathbf{x} \rangle^2}{\|\mathbf{w} - \mathbf{x}\|^2 + \epsilon}
+= \frac{\left(\sum_{i,j} w_{ij} x_{ij}\right)^2}{\sum_{i,j} (w_{ij} - x_{ij})^2 + \epsilon}
+$$
+
+Where:
+- $\mathbf{W}$ and $\mathbf{X}$ are local patches (e.g., kernel and input patch in convolution)
+- $w_{ij}$ and $x_{ij}$ are elements of the kernel and input patch, respectively
+- $\epsilon$ is a small constant for numerical stability
+
+This generalizes the Yat-product to convolutional (patch-wise) operations.
+
+
+## Supported Frameworks & API
+
+The `YatNMN` layer (for dense operations) and `YatConv` (for convolutional operations) are the core components. Below is a summary of their availability and features per framework:
+
+| Framework      | `YatNMN` Path                 | `YatConv` Path                | Core Layer | DropConnect | Ternary Network | Recurrent Layer |
+|----------------|-------------------------------|-------------------------------|------------|-------------|-----------------|-----------------|
+| **Flax (Linen)** | `src/nmn/linen/nmn.py`        | (Available)                   | ✅         |             |                 | 🚧              |
+| **Flax (NNX)**   | `src/nmn/nnx/nmn.py`          | `src/nmn/nnx/yatconv.py`      | ✅         | ✅          | 🚧              | 🚧              |
+| **Keras**      | `src/nmn/keras/nmn.py`        | (Available)                   | ✅         |             |                 | 🚧              |
+| **PyTorch**    | `src/nmn/torch/nmn.py`        | (Available)                   | ✅         |             |                 | 🚧              |
+| **TensorFlow** | `src/nmn/tf/nmn.py`           | (Available)                   | ✅         |             |                 | 🚧              |
+
+*Legend: ✅ Implemented, 🚧 To be implemented / In Progress, (Available) - Assumed available if NMN is, specific path might vary or be part of the NMN module.*
+
+## Installation
+
+```bash
+pip install nmn
+```
+
+## Usage Example (Flax NNX)
+
+```python
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from nmn.nnx.nmn import YatNMN
+from nmn.nnx.yatconv import YatConv
+
+# Example YatNMN (Dense Layer)
+model_key, param_key, drop_key, input_key = jax.random.split(jax.random.key(0), 4)
+in_features, out_features = 3, 4
+layer = YatNMN(in_features=in_features, out_features=out_features, rngs=nnx.Rngs(params=param_key, dropout=drop_key))
+dummy_input = jax.random.normal(input_key, (2, in_features)) # Batch size 2
+output = layer(dummy_input)
+print("YatNMN Output Shape:", output.shape)
+
+# Example YatConv (Convolutional Layer)
+conv_key, conv_param_key, conv_input_key = jax.random.split(jax.random.key(1), 3)
+in_channels, out_channels = 3, 8
+kernel_size = (3, 3)
+conv_layer = YatConv(
+    in_features=in_channels,
+    out_features=out_channels,
+    kernel_size=kernel_size,
+    rngs=nnx.Rngs(params=conv_param_key)
+)
+dummy_conv_input = jax.random.normal(conv_input_key, (1, 28, 28, in_channels)) # Batch 1, 28x28 image, in_channels
+conv_output = conv_layer(dummy_conv_input)
+print("YatConv Output Shape:", conv_output.shape)
+
+```
+*Note: Examples for other frameworks (Keras, PyTorch, TensorFlow, Flax Linen) can be found in their respective `nmn.<framework>` modules and upcoming documentation.*
+
+## Roadmap
+
+-   [ ] Implement recurrent layers (`YatRNN`, `YatLSTM`, `YatGRU`) for all supported frameworks.
+-   [ ] Develop Ternary Network versions of Yat layers for NNX.
+-   [ ] Add more comprehensive examples and benchmark scripts for various tasks (vision, language).
+-   [ ] Publish detailed documentation and API references.
+-   [ ] Conduct and publish thorough performance benchmarks against traditional layers.
+
+## Contributing
+
+Contributions are welcome! If you'd like to contribute, please feel free to:
+-   Open an issue on the [Bug Tracker](https://github.com/mlnomadpy/nmn/issues) to report bugs or suggest features.
+-   Submit a pull request with your improvements.
+-   Help expand the documentation or add more examples.
+
+## License
+
+This project is licensed under the **GNU Affero General Public License v3**. See the [LICENSE](LICENSE) file for details.
+
+## Citation
+
+If you use `nmn` in your research, please consider citing the original papers that inspired this work:
+
+> Deep Learning 2.0: Artificial Neurons that Matter: Reject Correlation - Embrace Orthogonality
+>
+> Deep Learning 2.1: Mind and Cosmos - Towards Cosmos-Inspired Interpretable Neural Networks
+
+A BibTeX entry will be provided once the accompanying paper for this library is published.
+
+## Citing
+
+If you use this work, please cite the paper:
+
+```bibtex
+@article{taha2024dl2,
+  author    = {Taha Bouhsine},
+  title     = {Deep Learning 2.0: Artificial Neurons that Matter: Reject Correlation - Embrace Orthogonality},
+}
+```
+
+
+```bibtex
+@article{taha2025dl2,
+  author    = {Taha Bouhsine},
+  title     = {Deep Learning 2.1: Mind and Cosmos - Towards Cosmos-Inspired Interpretable Neural Networks},
+}
+```

{nmn-0.1.4.dist-info → nmn-0.1.5.dist-info}/RECORD

@@ -1,14 +1,14 @@
 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=gWe8EL-aUm7be03M9O5R3XdBb92EpBEFsylrY6BA60c,4871
+nmn/nnx/nmn.py,sha256=tPNUtF8Lmv_B1TgMoVXfMQ9x0IPGKjSyAP6HnZ-YBsM,5651
 nmn/nnx/yatattention.py,sha256=chjtUKJtaR7ROPnNqkicbvMs7hzZKE0fIo_8cTNiju8,26601
-nmn/nnx/yatconv.py,sha256=xUH9NBY1fIDZeTA9GdgmqR_DJiQJgwU2uDrgxqirKmU,12308
+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.4.dist-info/METADATA,sha256=k28p055Dr6WWVQcb01uinFRiT5R-CAvdKz33fqZ85g4,5032
-nmn-0.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-nmn-0.1.4.dist-info/licenses/LICENSE,sha256=kbZSd5WewnN2PSjvAC6DprP7pXx6NUNsnltmU2Mz1yA,34519
-nmn-0.1.4.dist-info/RECORD,,
+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,,

nmn-0.1.4.dist-info/METADATA

@@ -1,119 +0,0 @@
-Metadata-Version: 2.4
-Name: nmn
-Version: 0.1.4
-Summary: a neuron that matter
-Project-URL: Homepage, https://github.com/mlnomadpy/nmn
-Project-URL: Bug Tracker, https://github.com/mlnomadpy/my_package/issues
-Author-email: Taha Bouhsine <yat@mlnomads.com>
-License-File: LICENSE
-Classifier: License :: OSI Approved :: GNU Affero General Public License v3
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-
-# nmn
-Not the neurons we want, but the neurons we need
-
-[![PyPI version](https://img.shields.io/pypi/v/nmn.svg)](https://pypi.org/project/nmn/)
-[![Downloads](https://static.pepy.tech/badge/nmn)](https://pepy.tech/project/nmn)
-[![Downloads/month](https://static.pepy.tech/badge/nmn/month)](https://pepy.tech/project/nmn)
-[![GitHub stars](https://img.shields.io/github/stars/mlnomadpy/nmn?style=social)](https://github.com/mlnomadpy/nmn)
-[![GitHub forks](https://img.shields.io/github/forks/mlnomadpy/nmn?style=social)](https://github.com/mlnomadpy/nmn)
-[![GitHub issues](https://img.shields.io/github/issues/mlnomadpy/nmn)](https://github.com/mlnomadpy/nmn/issues)
-[![PyPI - License](https://img.shields.io/pypi/l/nmn)](https://pypi.org/project/nmn/)
-[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/nmn)](https://pypi.org/project/nmn/)
-
-## Overview
-
-**nmn** provides neural network layers for multiple frameworks (Flax, NNX, Keras, PyTorch, TensorFlow) that do not require activation functions to learn non-linearity. The main goal is to enable deep learning architectures where the layer itself is inherently non-linear, inspired by the paper:
-
-> Deep Learning 2.0: Artificial Neurons that Matter: Reject Correlation - Embrace Orthogonality
-
-## Math
-
-Yat-Product: 
-$$
-ⵟ(\mathbf{w},\mathbf{x}) := \frac{\langle \mathbf{w}, \mathbf{x} \rangle^2}{\|\mathbf{w} - \mathbf{x}\|^2 + \epsilon} = \frac{ \|\mathbf{x}\|^2  \|\mathbf{w}\|^2 \cos^2 \theta}{\|\mathbf{w}\|^2 - 2\mathbf{w}^\top\mathbf{x} + \|\mathbf{x}\|^2 + \epsilon} = \frac{ \|\mathbf{x}\|^2  \|\mathbf{w}\|^2 \cos^2 \theta}{((\mathbf{x}-\mathbf{w})\cdot(\mathbf{x}-\mathbf{w}))^2 + \epsilon}.
-$$
-
-**Explanation:**
-- $\mathbf{w}$ is the weight vector, $\mathbf{x}$ is the input vector.
-- $\langle \mathbf{w}, \mathbf{x} \rangle$ is the dot product between $\mathbf{w}$ and $\mathbf{x}$.
-- $\|\mathbf{w} - \mathbf{x}\|^2$ is the squared Euclidean distance between $\mathbf{w}$ and $\mathbf{x}$.
-- $\epsilon$ is a small constant for numerical stability.
-- $\theta$ is the angle between $\mathbf{w}$ and $\mathbf{x}$.
-
-This operation:
-- **Numerator:** Squares the similarity (dot product) between $\mathbf{w}$ and $\mathbf{x}$, emphasizing strong alignments.
-- **Denominator:** Penalizes large distances, so the response is high only when $\mathbf{w}$ and $\mathbf{x}$ are both similar in direction and close in space.
-- **No activation needed:** The non-linearity is built into the operation itself, allowing the layer to learn complex, non-linear relationships without a separate activation function.
-- **Geometric view:** The output is maximized when $\mathbf{w}$ and $\mathbf{x}$ are both large in norm, closely aligned (small $\theta$), and close together in Euclidean space.
-
-Yat-Conv:
-$$
-ⵟ^*(\mathbf{W}, \mathbf{X}) := \frac{\langle \mathbf{w}, \mathbf{x} \rangle^2}{\|\mathbf{w} - \mathbf{x}\|^2 + \epsilon}
-= \frac{\left(\sum_{i,j} w_{ij} x_{ij}\right)^2}{\sum_{i,j} (w_{ij} - x_{ij})^2 + \epsilon}
-$$
-
-Where:
-- $\mathbf{W}$ and $\mathbf{X}$ are local patches (e.g., kernel and input patch in convolution)
-- $w_{ij}$ and $x_{ij}$ are elements of the kernel and input patch, respectively
-- $\epsilon$ is a small constant for numerical stability
-
-This generalizes the Yat-product to convolutional (patch-wise) operations.
-
-
-## Supported Frameworks & Tasks
-
-### Flax (JAX)
-- `YatNMN` layer implemented in `src/nmn/linen/nmn.py`
-- **Tasks:**
-  - [x] Core layer implementation
-  - [ ] Recurrent layer (to be implemented)
-
-### NNX (Flax NNX)
-- `YatNMN` layer implemented in `src/nmn/nnx/nmn.py`
-- **Tasks:**
-  - [x] Core layer implementation
-  - [ ] Recurrent layer (to be implemented)
-
-### Keras
-- `YatNMN` layer implemented in `src/nmn/keras/nmn.py`
-- **Tasks:**
-  - [x] Core layer implementation
-  - [ ] Recurrent layer (to be implemented)
-
-### PyTorch
-- `YatNMN` layer implemented in `src/nmn/torch/nmn.py`
-- **Tasks:**
-  - [x] Core layer implementation
-  - [ ] Recurrent layer (to be implemented)
-
-### TensorFlow
-- `YatNMN` layer implemented in `src/nmn/tf/nmn.py`
-- **Tasks:**
-  - [x] Core layer implementation
-  - [ ] Recurrent layer (to be implemented)
-
-## Installation
-
-```bash
-pip install nmn
-```
-
-## Usage Example (Flax)
-
-```python
-from nmn.nnx.nmn import YatNMN
-from nmn.nnx.yatconv import YatConv
-# ... use as a Flax module ...
-```
-
-## Roadmap
-- [ ] Implement recurrent layers for all frameworks
-- [ ] Add more examples and benchmarks
-- [ ] Improve documentation and API consistency
-
-## License
-GNU Affero General Public License v3