# Flax NNX Implementation of NMN

This directory contains the reference Flax NNX implementation of Neural Matter Network (NMN) components. Being the reference implementation, it contains the full suite of NMN features and algorithms.

## Status Comparison

This is the **Reference Implementation** ✅. All features are implemented here first.

### Implemented Features ✅
- **Core Layers**
  - `YatNMN` (Dense layer equivalent)
- **Convolutional Layers**
  - `YatConv` (Generic N-dimensional Convolution)
  - `YatConvTranspose` (Generic N-dimensional Transposed Convolution)
- **Recurrent Neural Networks (RNN)**
  - `YatGRU`
  - `YatLSTM`
  - `YatSimpleRNN`
- **Attention Mechanisms**
  - `YatAttention`
  - `MultiHeadYatAttention`
  - `RotaryYatAttention`
  - `SphericalYatPerformer`
- **Squashing Functions**
  - `SoftTanh`
  - `SofterSigmoid`
  - `Softermax`

## Usage Example

```python
from flax import nnx
from nmn.nnx import YatConv, YatNMN, SoftTanh

class Model(nnx.Module):
    def __init__(self, rngs: nnx.Rngs):
        self.conv1 = YatConv(1, 32, kernel_size=(3, 3), rngs=rngs)
        self.conv2 = YatConv(32, 64, kernel_size=(3, 3), rngs=rngs)
        self.linear = YatNMN(64, 10, rngs=rngs)
        self.activation = SoftTanh()

    def __call__(self, x):
        x = self.activation(self.conv1(x))
        x = self.activation(self.conv2(x))
        x = x.mean(axis=(1, 2)) # Global Average Pooling
        x = self.linear(x)
        return nnx.softmax(x)
```
