ModelStudio

ModelStudio is an early-stage AI tensor framework with a working CPU MVP, 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.

Installation

From PyPI:

python -m pip install modelstudio

For development:

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

Current Features

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.

Test

python -m pytest

Smoke test:

python scripts/smoke_test.py

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

Backend Status

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.

Documentation

• Backend architecture
• Autograd design
• Releasing
• Contributing

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.