ModelStudio

ModelStudio is an early-stage AI tensor framework. Version 0.6.0 provides a CPU tensor/autograd MVP with neural-network modules, optimizers, serialization, data loading, graph tracing metadata, backend status inspection, a public CUDA availability namespace, and small LLM-oriented building blocks.

It is not a PyTorch or TensorFlow replacement. The default PyPI package is CPU-only. CUDA, ROCm, and oneAPI remain explicit scaffolds until real kernels are built and tested in hardware-backed environments.

Installation

From PyPI:

python -m pip install modelstudio

For development:

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

Feature Table

Architecture

Python frontend
  -> Tensor, nn, optim, autograd, ops
  -> runtime dispatcher
  -> backend interface
  -> NumPy CPU backend today
  -> optional native CPU / CUDA / ROCm / oneAPI extensions later

Native scaffold
  -> core metadata
  -> dispatcher interfaces
  -> CPU kernel prototypes
  -> CUDA, ROCm, oneAPI backend directories

Backend Status

import modelstudio as ms

print(ms.backends.status())
print(ms.backends.native_cpu_available())

Expected shape:

{
    "cpu": {"available": True, "native": False},
    "cuda": {"available": False, "built": False, "device_count": 0, "reason": "..."},
    "rocm": {"available": False, "reason": "..."},
    "oneapi": {"available": False, "reason": "..."},
}

The production CPU path is the NumPy backend. ms.backends.use_native_cpu(True) raises ModelStudioBackendUnavailable unless a future optional native extension is actually installed. Unsupported accelerator devices fail with ModelStudioBackendUnavailable.

CUDA availability can also be checked through the public namespace:

print(ms.cuda.is_available())
print(ms.cuda.device_count())
print(ms.cuda.memory_summary())

In the CPU-only wheel, explicit CUDA tensor requests raise a clear runtime error instead of falling back to CPU.

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)

Functional API

import modelstudio as ms
from modelstudio import nn
from modelstudio.nn import functional as F

model = nn.Linear(4, 2)
x = ms.random.randn((8, 4))
target = ms.random.randn((8, 2))
loss = F.mse_loss(F.relu(F.linear(x, model.weight, model.bias)), target)

Tracing

import modelstudio as ms
from modelstudio.nn import functional as F

x = ms.random.randn((4, 3))
w = ms.random.randn((3, 2))
graph = ms.trace(lambda a, b: F.relu(a @ b), x, w)
print(graph)

Tracing captures operation names and tensor metadata. It does not optimize or execute graphs yet. ms.compile(fn) remains a documented no-op that returns the original callable.

Random And Linalg

ms.random.seed(123)
x = ms.random.normal((4, 3), mean=0.0, std=1.0)
w = ms.random.uniform((3, 2), low=-0.1, high=0.1)
y = ms.linalg.matmul(x, w)
print(ms.linalg.norm(y).item())

Comparisons

x = ms.tensor([1.0, 2.0, 3.0])
y = ms.tensor([1.0, 2.1, 3.0])
print(ms.isclose(x, y, atol=0.05))
print(ms.allclose(x, y, atol=0.05))
print((x > 1.5).any().item())

Comparison and logical outputs are bool tensors and do not track gradients.

Checkpointing

model = nn.Linear(4, 2)
optimizer = ms.optim.AdamW(model.parameters(), lr=1e-3)
ms.save_checkpoint("checkpoint.ms", model=model, optimizer=optimizer, extra={"epoch": 1})
checkpoint = ms.load_checkpoint("checkpoint.ms", model=model, optimizer=optimizer, map_location="cpu")

Checkpoint loading validates structure and model state. CPU is the only accepted map_location in the current release.

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 examples/backend_status.py
python examples/tracing_demo.py
python examples/functional_training.py
python examples/random_linalg_demo.py
python examples/cuda_tensor_demo.py
python examples/cuda_mlp_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
python benchmarks/bench_elementwise.py
python benchmarks/bench_trace.py
python benchmarks/bench_cuda_elementwise.py
python benchmarks/bench_cuda_matmul.py

Documentation

• Backend status
• CUDA status
• Tracing
• Functional API
• Random namespace
• Linalg namespace
• Comparison ops
• Tensor API
• Neural network API
• Data utilities
• Training
• Modules
• Serialization
• 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 optional native CPU kernels only after a safe Python extension exists.
• Build a real optional CUDA package after tensor storage, kernels, bindings, and hardware-backed CI are in place.
• Add tested ROCm and oneAPI packages after CUDA establishes the accelerator backend contract.
• Improve compiler graph capture, analysis passes, and lowering.