ModelStudio

ModelStudio is an early-stage AI tensor framework. Version 0.2.0 provides a CPU tensor/autograd MVP with neural-network modules, optimizers, serialization, basic data loading, and small LLM-oriented building blocks.

It is not a PyTorch or TensorFlow replacement. CPU is the only working backend. CUDA, ROCm, and oneAPI remain explicit scaffolds until real kernels are built and tested.

Installation

From PyPI:

python -m pip install modelstudio

For development:

python -m pip install -e ".[dev]"

Feature Table

Backend Status

Unsupported accelerator devices fail with ModelStudioBackendUnavailable.

Tensor Example

import modelstudio as ms

x = ms.randn((32, 784), requires_grad=True)
w = ms.randn((784, 10), requires_grad=True)
loss = (x @ w).mean()
loss.backward()
print(w.grad)

MLP Example

import modelstudio as ms
from modelstudio import nn


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        return self.fc2(ms.gelu(self.fc1(x)))


model = MLP()
optimizer = ms.optim.AdamW(model.parameters(), lr=3e-4)
x = ms.randn((16, 784))
target = ms.randn((16, 10))
loss = ms.mse_loss(model(x), target)
optimizer.zero_grad()
loss.backward()
optimizer.step()

State Dict and Save/Load

model = nn.Linear(4, 2)
ms.save(model.state_dict(), "model.ms")
state = ms.load("model.ms")
model.load_state_dict(state)

DataLoader

from modelstudio import data

dataset = data.TensorDataset(ms.randn((8, 4)), ms.arange(8))
loader = data.DataLoader(dataset, batch_size=2, shuffle=False)
for xb, yb in loader:
    print(xb.shape, yb.shape)

Embedding

emb = nn.Embedding(num_embeddings=100, embedding_dim=32)
tokens = ms.tensor([[1, 2, 3]], dtype=ms.int64)
print(emb(tokens).shape)

Cross Entropy

logits = ms.randn((4, 10), requires_grad=True)
targets = ms.tensor([1, 2, 3, 4], dtype=ms.int64)
loss = ms.cross_entropy(logits, targets)
loss.backward()

TransformerBlock

block = nn.TransformerBlock(embed_dim=16, num_heads=4)
x = ms.randn((2, 8, 16), requires_grad=True)
y = block(x)
print(y.shape)

Commands

python -m pytest
python scripts/smoke_test.py
python examples/train_mlp.py
python examples/train_classifier.py
python examples/tiny_transformer.py
python examples/save_load.py
python benchmarks/bench_matmul.py
python benchmarks/bench_mlp.py
python benchmarks/bench_attention.py
python benchmarks/bench_dataloader.py

Documentation

• Tensor API
• Neural network API
• Data utilities
• Backend architecture
• Autograd design
• Releasing
• Contributing

Roadmap

• Expand tensor and autograd coverage.
• Wire native CPU kernels into Python bindings.
• Add tested CUDA, ROCm, and oneAPI packages when hardware-backed CI exists.
• Improve compiler graph capture and lowering.