# Linen Implementation of NMN

This directory contains the Flax Linen (JAX) implementation of Neural Matter Network (NMN) components.

## Status Comparison

Compared to the reference `nnx` implementation, the Linen backend implements the core dense and convolutional layers but currently lacks transposed convolutions, attention mechanisms, and RNNs.

### Implemented Features ✅
- **Core Layers**
  - `YatNMN` (Dense layer equivalent)
- **Convolutional Layers**
  - `YatConv1D`
  - `YatConv2D`
  - `YatConv3D`

### Missing Features 🚧
The following components from the reference implementation are currently missing or work-in-progress:

- **Transposed Convolution**
  - `YatConvTranspose1D`
  - `YatConvTranspose2D`
  - `YatConvTranspose3D`
- **Recurrent Neural Networks (RNN)**
  - `YatGRU`
  - `YatLSTM`
  - `YatSimpleRNN`
- **Attention Mechanisms**
  - `YatAttention`
  - `MultiHeadYatAttention`
  - `RotaryYatAttention`
  - `SphericalYatPerformer`
- **Squashing Functions**
  - `SoftTanh`
  - `SofterSigmoid`
  - `Softermax`

## Usage Example

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
from nmn.linen import YatConv2D, YatNMN

class Model(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = YatConv2D(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = YatConv2D(features=64, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))  # flatten
        x = YatNMN(features=10)(x)
        return nn.softmax(x)
```
