Migrating from PyTorch to SNEPPX-Algo

A side-by-side mapping from common torch APIs to SneppX_ALG equivalents. The SNEPPX Python layer mirrors PyTorch's ergonomics but requires the C backend (_HAS_C_BACKEND is True) for training; pure-NumPy fallbacks exist for inference-grade ops.

Installation

# PyTorch
pip install torch
# SneppX
pip install sneppx-alg            # wheels include the C backend
# or from source:
git clone https://github.com/ammar49-cyber/sneppx-alg
cmake -B build -G Ninja -DCMAKE_BUILD_TYPE=Release
cmake --build build --config Release

Set the path (from source builds):

$env:PYTHONPATH = "bindings/python"

Tensor

PyTorch	SNEPPX-Algo
torch.tensor([...])	Tensor([...])
torch.zeros(4, 8)	Tensor.zeros(4, 8)
torch.randn(4, 8)	Tensor.randn(4, 8)
x.numpy()	x.numpy()
x.to("cuda")	x.to("cuda")
x @ y / x.matmul(y)	x @ y (__matmul__)
x.requires_grad	x.requires_grad
x.grad	x.grad

# PyTorch
import torch
x = torch.randn(4, 8)
w = torch.randn(8, 16, requires_grad=True)
y = x @ w
y.sum().backward()

# SNEPPX
from SneppX_ALG import Tensor
x = Tensor.randn((4, 8))
w = Tensor.randn((8, 16), requires_grad=True)
y = x @ w
y.backward()           # autodiff tape; requires C backend

nn.Module ↔ Module

PyTorch	SNEPPX-Algo
nn.Module	Module
nn.Linear	Linear
nn.Embedding	Embedding
nn.LayerNorm	LayerNorm
nn.RMSNorm	RMSNorm
nn.Dropout	Dropout
nn.Sequential	Sequential
nn.MultiheadAttention	MultiheadAttention
nn.TransformerEncoderLayer	TransformerBlock
nn.Transformer	Transformer
model.parameters()	model.parameters()
model.state_dict()	model.state_dict()
model.load_state_dict(...)	model.load_state_dict(...)
model.to(device)	model.to(device)
model.train() / model.eval()	model.train() / model.eval()

# PyTorch
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(784, 256), nn.GELU(), nn.Linear(256, 10))
    def forward(self, x):
        return self.net(x.flatten(1))

# SNEPPX
from SneppX_ALG import Module, Linear, Sequential, TransformerBlock
class MLP(Module):
    def __init__(self):
        super().__init__()
        self.net = Sequential(Linear(784, 256), GELU(), Linear(256, 10))
    def forward(self, x):
        return self.net(x.reshape((-1, 784)))

Optimizers

PyTorch	SNEPPX-Algo
torch.optim.SGD	SGD
torch.optim.AdamW	AdamW
(Lion)	Lion
(LAMB)	LAMB
torch.optim.lr_scheduler.CosineAnnealingLR	CosineAnnealingLR

# PyTorch
opt = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=100)

# SNEPPX
from SneppX_ALG import AdamW, CosineAnnealingLR
opt   = AdamW(model.parameters(), lr=2e-4, weight_decay=0.01)
sched = CosineAnnealingLR(opt, min_lr=1e-5, max_lr=2e-4, total_steps=100)

Training loop

PyTorch	SNEPPX-Algo
loss.backward()	loss.backward()
optimizer.step()	optimizer.step()
optimizer.zero_grad()	optimizer.zero_grad()
(Trainer)	Trainer.fit(loader)

# SNEPPX — tape-based autodiff via the C backend
from SneppX_ALG import Tensor, AdamW, MSELoss
opt = AdamW(model.parameters(), lr=1e-3)
loss_fn = MSELoss()
for x, y in loader:
    opt.zero_grad()
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    opt.step()

The SneppX_ALG.Trainer class wraps the C training loop (Trainer.fit); for low-level control use the optimizer directly as above.

Data

PyTorch	SNEPPX-Algo
torch.utils.data.Dataset	Dataset
torch.utils.data.TensorDataset	TensorDataset
torch.utils.data.DataLoader	DataLoader (interfacebindings.dataloader)
torch.utils.data.distributed.DistributedSampler	DistributedSampler

Distributed

PyTorch	SNEPPX-Algo
torch.distributed.init_process_group	init_process_group
torch.nn.parallel.DistributedDataParallel	DistributedDataParallel / DistributedWrapper
torch.distributed.launch / torchrun	launch(train_fn, num_gpus=...)
torch.distributed.is_initialized	DistributedContext.initialized

Generation

PyTorch (HF model.generate)	SNEPPX-Algo
GenerationConfig(...)	GenerationConfig(...)
model.generate(...)	generate(model, input_ids, ...)
LogitsWarperList	top_k_top_p_filtering(...)
AutoTokenizer	Tokenizer / SimpleTokenizer

from SneppX_ALG.interface_bindings.generation import generate, GenerationConfig
from SneppX_ALG import Tokenizer

gen_config = GenerationConfig(max_new_tokens=64, temperature=0.7, top_p=0.9)
tok = Tokenizer(vocab_size=32000)
ids = tok.encode("Hello, SneppX")
result = generate(model, ids, generation_config=gen_config)
print(tok.decode(result["output_ids"].tolist()[0]))

Quick reference table

Concept	torch	SneppX-Algo
Backend	CUDA	C11/C++20 (_SNEPPX_c) + CUDA (_HAS_CUDA)
Grad	autograd	tape-based backward()
RNG	torch.manual_seed	seed= on config
Device	torch.device	"cpu" / "cuda" strings
Dtype	torch.float32	"float32", Dtype.FLOAT32
Save	torch.save	model.save_checkpoint(path)

Gotchas

• SNEPPX Tensor uses 4-space indentation in Python, SNEPPX_ prefix on C APIs, and void parameter lists in C — not a Python concern, but the binding docstrings follow the same conventions.
• Linear.forward does x @ weight.T (matches PyTorch convention).
• Without the C backend, backward()/Trainer.fit raise RuntimeError. Build the extension (cmake --build build --config Release) to enable them.