ModelStudio

ModelStudio is an early, production-oriented foundation for a Python AI tensor framework. The current MVP provides a working CPU tensor API, NumPy-backed CPU kernels, reverse-mode autograd, neural-network modules, optimizers, dispatcher boundaries, and native backend scaffolding for future C++/CUDA/ROCm/oneAPI work.

It is not a complete PyTorch or TensorFlow replacement. GPU folders are intentionally scaffolding until real backend implementations are written, built, and tested.

What Works Today

• Python package modelstudio, commonly imported as ms
• CPU tensors with shape, dtype, device, strides, .numpy(), .item(), .to(), .detach(), .zero_grad(), and .backward()
• DTypes: float32, float64, int32, int64, bool
• Device parsing for cpu, cuda, rocm, and oneapi
• CPU creation ops: tensor, empty, zeros, ones, randn, arange
• CPU math ops: +, -, *, /, negation, @, sum, mean, reshape, transpose, .T, relu, gelu
• Reverse-mode autograd for the MVP ops, including broadcasting gradients
• no_grad() and is_grad_enabled()
• nn.Module, Parameter, Linear, ReLU, GELU, MSELoss, LayerNorm
• optim.SGD and optim.AdamW
• Dispatcher with CPU backend registered by default
• Unavailable CUDA, ROCm, and oneAPI backends that fail with explicit runtime errors

Not Implemented Yet

• Real CUDA kernels or memory management
• Real ROCm/HIP kernels or memory management
• Real oneAPI/SYCL kernels or memory management
• Native Python extension bindings
• Full compiler graph capture and lowering
• Serialization, distributed training, mixed precision, convolution, sparse tensors, and production profiler hooks

Unsupported GPU devices fail clearly, for example:

ModelStudioBackendUnavailable: CUDA backend is not built. Install modelstudio-cuda or build with MODELSTUDIO_ENABLE_CUDA=ON.

Architecture

Python frontend API
  |
  v
Tensor + Autograd
  |
  v
Ops modules
  |
  v
Runtime Dispatcher
  |
  +--> CPU Backend (NumPy MVP today, replaceable with C++ kernels)
  |
  +--> CUDA Backend placeholder (not built)
  |
  +--> ROCm Backend placeholder (not built)
  |
  +--> oneAPI Backend placeholder (not built)

Native scaffolding
  |
  +--> core tensor metadata, dtype, device, storage
  +--> dispatcher/backend interfaces
  +--> CPU kernel source layout
  +--> CUDA/ROCm/oneAPI source layout

The public API is intentionally independent from NumPy. NumPy is an internal CPU backend detail for the MVP, so later native CPU kernels can replace it without breaking Python users.

Install

From the repository root:

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

Test

python -m pytest

Example

import modelstudio as ms
from modelstudio import nn

x = ms.randn((32, 784), device="cpu", requires_grad=True)
w = ms.randn((784, 10), device="cpu", requires_grad=True)

y = x @ w
loss = y.mean()
loss.backward()

print(w.grad)

Training 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):
        x = ms.gelu(self.fc1(x))
        return self.fc2(x)


model = MLP()
optimizer = ms.optim.AdamW(model.parameters(), lr=3e-4)

x = ms.randn((16, 784))
target = ms.randn((16, 10))

pred = model(x)
loss = ms.mse_loss(pred, target)

optimizer.zero_grad()
loss.backward()
optimizer.step()

Run the included example:

python examples/train_mlp.py

Benchmarks

python benchmarks/bench_matmul.py
python benchmarks/bench_mlp.py

The CPU MVP routes through NumPy, so benchmarks are mainly smoke tests and baseline measurements for future native kernels.

Native Backend Strategy

Native source lives under csrc/ and uses the top-level C++ namespace modelstudio.

Build options:

option(MODELSTUDIO_ENABLE_CUDA "Build CUDA backend" OFF)
option(MODELSTUDIO_ENABLE_ROCM "Build ROCm backend" OFF)
option(MODELSTUDIO_ENABLE_ONEAPI "Build oneAPI backend" OFF)

The intended path is:

1. Keep the Python dispatcher contract stable.
2. Replace NumPy CPU kernels with native CPU kernels behind the same backend interface.
3. Add allocator and kernel implementations for CUDA, ROCm, and oneAPI in their backend folders.
4. Package GPU backends as optional install targets, such as modelstudio-cuda.
5. Require backend-specific tests before claiming support.

Roadmap

• Expand tensor op coverage and gradient coverage.
• Wire native CPU kernels into Python bindings.
• Add graph capture, shape inference, fusion, and lowering.
• Implement serialization and state dicts.
• Implement real CUDA, ROCm, and oneAPI backends with CI coverage on matching hardware.
• Add documentation pages beyond the README as APIs stabilize.