---

SQRTAI 🧠

A PyTorch neural network trained to approximate the complex square root of real numbers.

The model receives a real number and predicts the real and imaginary components of its square root.

🔬 Architecture & Features

SQRTAI is a fully connected neural network built with PyTorch.

• Input: 1 real value
• Output: 2 values [real, imag]
• Activation: GELU
• Loss: SmoothL1Loss
• Optimizer: AdamW
• Learning rate: 0.001
• Batch size: 32
• Training epochs: 300
• Input normalization: none

Architecture

Input (1)
   ↓
Linear(1 → 750)
   ↓
GELU
   ↓
Linear(750 → 1500)
   ↓
GELU
   ↓
Linear(1500 → 750)
   ↓
GELU
   ↓
Linear(750 → 2)
   ↓
[Real, Imaginary]

The model approximates:

√x = real + imag·i

Examples:

√100   ≈ 10 + 0i
√25    ≈ 5 + 0i
√-100  ≈ 0 + 10i
√-25   ≈ 0 + 5i

📚 Training Data

The model was trained on integer values in the range:

[-2500, 2499]

The targets were generated using PyTorch complex numbers:

complex_x = x.to(torch.complex64)
complex_y = torch.sqrt(complex_x)

The real and imaginary parts were then stored as two separate target values.

📊 Training Results

Example predictions from the trained model:

Input	SQRTAI	True
-1000	0.0013 + 31.6555i	0 + 31.6228i
-100	0.0014 + 9.9772i	0 + 10.0000i
-25	0.0067 + 5.1357i	0 + 5.0000i
-4	-0.0819 + 2.6917i	0 + 2.0000i
-1	0.3614 + 1.2544i	0 + 1.0000i
0	0.5484 + 0.8187i	0 + 0i
1	0.9438 + 0.4106i	1.0000 + 0i
4	2.1236 - 0.0022i	2.0000 + 0i
25	4.9417 - 0.0021i	5.0000 + 0i
100	10.0111 + 0.0040i	10.0000 + 0i
1000	31.7295 - 0.0045i	31.6228 + 0i

⚠️ Known Limitations

The model is an approximation, not an exact mathematical calculator.

The largest errors in the current model appear around the origin, especially at:

x = -1
x = 0
x = 1

For example:

x = -1
AI   → 0.3614 + 1.2544i
True → 0 + 1i

x = 0
AI   → 0.5484 + 0.8187i
True → 0 + 0i

x = 1
AI   → 0.9438 + 0.4106i
True → 1 + 0i

This means SQRTAI does not accurately reproduce the square-root function near x = 0.

For larger absolute input values, the approximation becomes significantly closer to the mathematical result.

🧮 Model Statistics

• Architecture: 1 → 750 → 1500 → 750 → 2
• Total parameters: 2,255,252
• Trainable parameters: 2,255,252
• Non-trainable parameters: 0
• Approximate FP32 size: ~8.6 MiB

💻 How to Use

import torch as t
import torch.nn as nn


class SQRTAI(nn.Module):
    def __init__(self):
        super().__init__()

        self.l1 = nn.Linear(1, 750)
        self.l2 = nn.Linear(750, 1500)
        self.l3 = nn.Linear(1500, 750)
        self.l4 = nn.Linear(750, 2)

        self.gelu = nn.GELU()

    def forward(self, x):
        x = self.gelu(self.l1(x))
        x = self.gelu(self.l2(x))
        x = self.gelu(self.l3(x))
        x = self.l4(x)

        return x


device = t.device("cuda" if t.cuda.is_available() else "cpu")

model = SQRTAI().to(device)

model.load_state_dict(
    t.load("SQRTAI.pth", map_location=device)
)

model.eval()


x_value = 100.0

with t.no_grad():
    x = t.tensor(
        [[x_value]],
        dtype=t.float32
    ).to(device)

    prediction = model(x)

real = prediction[0, 0].item()
imag = prediction[0, 1].item()

print(f"SQRTAI: {real:.4f} + {imag:.4f}i")

🧪 Example

Input:
100

SQRTAI:
10.0111 + 0.0040i

True:
10.0000 + 0i

📦 Model File

The trained weights are stored in:

SQRTAI.pth

🚀 Future Improvements

Possible future versions could improve accuracy around x ≈ 0, increase the training dataset, and investigate alternative architectures or training strategies.

---

In this work, I am not attempting to replace a calculator. This is a research project aimed at investigating whether a standard MLP can solve a problem involving complex numbers in the context of square roots.

Made with PyTorch 🔥