# PyTorch Implementation of NMN

This directory contains the PyTorch implementation of Neural Matter Network (NMN) components.

## Status Comparison

Compared to the reference `nnx` implementation, the PyTorch backend implements the core dense and convolutional layers, as well as transposed convolutions and some attention mechanisms. It is one of the more complete backends but still lacks RNNs and some advanced attention features.

### Implemented Features ✅
- **Core Layers**
  - `YatNMN` (Dense layer equivalent)
- **Convolutional Layers**
  - `YatConv1d`
  - `YatConv2d`
  - `YatConv3d`
- **Transposed Convolutional Layers**
  - `YatConvTranspose1d`
  - `YatConvTranspose2d`
  - `YatConvTranspose3d`
- **Attention Mechanisms**
  - `YatAttention`
  - `MultiHeadYatAttention`

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

- **Recurrent Neural Networks (RNN)**
  - `YatGRU`
  - `YatLSTM`
  - `YatSimpleRNN`
- **Advanced Attention Features**
  - `RotaryYatAttention`
  - `SphericalYatPerformer`
- **Squashing Functions**
  - `SoftTanh`
  - `SofterSigmoid`
  - `Softermax`

## Usage Example

```python
import torch
import torch.nn as nn
from nmn.torch import YatConv2d, YatNMN

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = YatConv2d(1, 32, kernel_size=3, stride=2, padding=1)
        self.relu1 = nn.ReLU()
        self.conv2 = YatConv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.relu2 = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = YatNMN(64, 10)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.relu1(self.conv1(x))
        x = self.relu2(self.conv2(x))
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return self.softmax(x)
```
