ModelStudio

ModelStudio is an early-stage AI tensor framework. Version 0.4.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)

0.4.0 Training Utilities

ms.manual_seed(123)
model = nn.Linear(4, 2)
optimizer = ms.optim.AdamW(model.parameters(), lr=1e-3)
state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
ms.save(state, "checkpoint.ms")

New CPU-only helpers include ms.concat, ms.stack, Tensor.flatten, Tensor.squeeze, Tensor.unsqueeze, nn.init, nn.Dropout, nn.BatchNorm1d, nn.Conv1d, nn.Conv2d, nn.AvgPool2d, nn.MaxPool2d, and nn.utils gradient clipping.

NumPy Interop

x = ms.asarray([[1, 2, 3], [4, 5, 6]], dtype=ms.float32)
arr = ms.to_numpy(x)
y = ms.from_numpy(arr)

CPU uses NumPy internally. Normal examples prefer ModelStudio APIs; ms.numpy is exposed for advanced users who explicitly want NumPy access.

Schedulers and Metrics

optimizer = ms.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = ms.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5)
scheduler.step()

acc = ms.metrics.accuracy(logits, targets)

Checkpointing

ms.save_checkpoint("checkpoint.ms", model=model, optimizer=optimizer, scheduler=scheduler, extra={"epoch": 1})
checkpoint = ms.load_checkpoint("checkpoint.ms", model=model, optimizer=optimizer, scheduler=scheduler)

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 examples/train_cnn_toy.py
python examples/dropout_batchnorm.py
python examples/checkpoint_training.py
python examples/numpy_interop.py
python examples/scheduler_training.py
python examples/checkpoint_resume.py
python examples/metrics_demo.py
python benchmarks/bench_matmul.py
python benchmarks/bench_mlp.py
python benchmarks/bench_attention.py
python benchmarks/bench_dataloader.py
python benchmarks/bench_conv.py
python benchmarks/bench_dropout.py
python benchmarks/bench_creation.py
python benchmarks/bench_manipulation.py

Documentation

• Tensor API
• Neural network API
• Data utilities
• Training
• Modules
• Serialization
• Randomness
• Native backend roadmap
• NumPy interop
• Tensor creation
• Tensor manipulation
• Optimizers
• Checkpointing
• Metrics
• 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.