TensorStudio

TensorStudio is a compact C++ tensor and autograd engine with a Python API for learning, experimentation, and lightweight ML workloads.

TensorStudio 2.1.0 is a CPU-only stable API foundation with native C++ threading, storage reuse, SIMD-friendly typed kernels, and optional CBLAS/Accelerate matrix multiplication when available. It adds native stable softmax/logsumexp, batched matrix multiplication, statistical reductions, boolean reductions, seeded random distributions, and a hardened eager autograd lifecycle. The neural-network layer now includes grouped/depthwise/1D/ transposed convolution, normalization layers, embeddings, richer activations, initializers, additional losses, and model summaries. The project layer adds dataset factories, deterministic train/validation splitting, metrics, callbacks, multi-format configs, checkpoint resume helpers, and starter project templates. The interchange layer adds richer NPZ metadata, optional SafeTensors weight storage, ONNX metadata inspection, supported-subset ONNX import/execution, and model metadata helpers. The hardware layer now exposes formal CPU/CUDA/Metal device descriptors, backend availability reporting, device-aware storage boundaries, explicit transfer APIs, and backend benchmark artifacts while keeping the published wheels honest about CPU-only execution. The graph layer adds constrained symbolic tracing, JSON graph serialization, simple graph optimization, eager-backed graph execution, profiling hooks, and memory-plan metadata. The ecosystem layer adds experimental COO sparse tensors, public dataset format readers, tiny model-zoo factories, language-model helpers, quantization research utilities, a custom-kernel registry, single-process distributed planning helpers, and an optional ONNX Runtime adapter. The v2 foundation adds dataset manifests/checksums, a small dataset cache wrapper, CSR sparse matrices, compact attention/Transformer encoder blocks, quantization calibration, CPU DLPack import, and ONNX Runtime provider diagnostics. Version 2.1.0 extends the hardware boundary with allocator/runtime/logical-device metadata, kernel placement and transfer planning, descriptor-only kernel manifests, native storage telemetry, ONNX Runtime named-input inference, safe model-format inspection for Keras/SavedModel/HDF5/TFLite files, and deterministic image-folder manifests. It is eager-first, intentionally compact, and not a replacement for mature ML frameworks.

Install

From PyPI:

python -m pip install tensorstudio

From a source checkout:

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

Install optional extras for ONNX export, optional ONNX Runtime delegation, and Pillow-backed image inputs:

python -m pip install "tensorstudio[onnx,vision]"
python -m pip install "tensorstudio[onnxruntime]"

Build source and wheel distributions:

python -m build
python -m twine check dist/*

End users should install wheels and should not need CMake. Source builds require a C++20 compiler because the native extension is implemented in C++.

Platform Setup

Windows is the primary release target. Use Python from python.org and install Microsoft C++ Build Tools or Visual Studio with the Desktop development with C++ workload before building from source:

python -m pip install -U pip
python -m pip install -e ".[dev]"
pytest -q

Linux source builds need GCC or Clang, CMake, and Python development headers:

python -m pip install -U pip
python -m pip install -e ".[dev]"
pytest -q

macOS source builds need Xcode Command Line Tools:

xcode-select --install
python -m pip install -U pip
python -m pip install -e ".[dev]"
pytest -q

Quickstart

import tensorstudio as ts

x = ts.tensor([[1.0, 2.0], [3.0, 4.0]])
y = ts.ones((2, 2))

print((x + y).tolist())
print((x @ y).numpy())
print(x.reshape((4,)).tolist())
print(x[0, :].tolist())
print(x.unsqueeze(0).permute(1, 2, 0).shape)

Tensor API

TensorStudio supports CPU tensors with float32, float64, int32, int64, and bool dtypes.

import tensorstudio as ts

ts.manual_seed(7)

a = ts.zeros((2, 3))
b = ts.rand((2, 3))
c = ts.eye(3)
d = ts.linspace(0.0, 1.0, 5)
labels = ts.randint((4,), low=0, high=3, seed=3)
mask = ts.bernoulli((2, 3), probability=0.25, seed=5)

print(a.shape, a.strides, a.device, a.is_contiguous)
print((b.clamp(0.2, 0.8) + 1).mean().item())
print(b.sum(axis=1).tolist())
print(ts.concat([b, b], axis=0).shape, b.astype("float64").dtype)
print(c.tolist(), d.tolist())
print(labels.tolist(), mask.any(axis=1).tolist())
print(ts.zeros_like(b).shape, ts.randn_like(b, seed=11).dtype)

Arithmetic promotion is explicit and inspectable:

print(ts.promote_types("int32", "float32"))        # float32
print(ts.result_type("int64", "int32", op="div")) # float32
print(ts.result_type("int64", "float32", op="gt")) # bool

Advanced Math

Native C++ elementwise math includes trigonometric functions and numerically useful helpers with autograd support:

import tensorstudio as ts

x = ts.tensor([0.1, 0.2, 0.3], requires_grad=True)
y = ts.sin(x) + x.cos() + x.log1p() + x.rsqrt()
loss = y.mean()
loss.backward()

print(loss.item())
print(x.grad.tolist())

Stable reductions and normalized probabilities are available as Tensor methods, functional ops, and tensorstudio.math helpers:

values = ts.tensor([[1000.0, 1001.0, 999.0], [1.0, 2.0, 3.0]])

print(ts.math.variance(values).item())
print(ts.math.std(values, axis=0).tolist())
print(ts.math.norm(values, ord=2).item())
print(values.softmax(axis=1).tolist())
print(ts.logsumexp(values, axis=1).tolist())

Batched matrix multiplication and a small documented einsum subset cover common model and scientific-programming patterns:

left = ts.randn((2, 3, 4), seed=1)
right = ts.randn((2, 4, 5), seed=2)

print((left @ right).shape)
print(ts.bmm(left, right).shape)
print(ts.einsum("bij,bjk->bik", left, right).shape)

Autograd

import tensorstudio as ts

x = ts.tensor([1.0, 2.0, 3.0], requires_grad=True)
loss = (x * x).mean()
loss.backward()

print(x.grad.tolist())

Reuse a graph explicitly when needed:

loss = (x * x).sum()
loss.backward(retain_graph=True)
loss.backward()

Use no_grad() when you want eager computation without recording a graph:

with ts.no_grad():
    y = x * 2
    x.zero_()

Neural Networks

import tensorstudio as ts
from tensorstudio import nn, optim

ts.manual_seed(0)

model = nn.Sequential(nn.Linear(1, 8), nn.Tanh(), nn.Linear(8, 1))
optimizer = optim.SGD(model.parameters(), lr=0.05, momentum=0.9)
scheduler = optim.StepLR(optimizer, step_size=50, gamma=0.5)
criterion = nn.MSELoss()

x = ts.tensor([[0.0], [1.0], [2.0], [3.0]])
y = ts.tensor([[1.0], [3.0], [5.0], [7.0]])

for _ in range(100):
    optimizer.zero_grad()
    loss = criterion(model(x), y)
    loss.backward()
    optim.clip_grad_norm_(model.parameters(), max_norm=10.0)
    optimizer.step()
    scheduler.step()

print(loss.item())
print(model.state_dict().keys())
print(model.parameter_count())

The neural-network surface also includes initialization helpers, normalization layers, embeddings, grouped/depthwise/1D/transposed convolution, adaptive/global pooling, richer activations, and model summaries:

model = nn.Sequential(
    nn.Conv2d(1, 8, kernel_size=3, padding=1),
    nn.BatchNorm2d(8),
    nn.GELU(),
    nn.GlobalAvgPool2d(),
    nn.Flatten(),
    nn.Linear(8, 10),
)
nn.init.kaiming_uniform_(model[0].weight, nonlinearity="relu", seed=7)
print(nn.summary(model, input_shape=(1, 1, 28, 28))["total_parameters"])

Vision

TensorStudio includes a practical computer-vision namespace for local image classification workflows: Pillow-backed image IO, transform pipelines, deterministic augmentations, ImageFolder datasets, metrics, image grids, bounding-box drawing, and compact CNN classifiers running through native Conv2d/pooling kernels. The current vision surface includes batch-aware transforms, color jitter, random resized crop, rotation, affine transforms, cutout, mixup, CutMix, detection helpers, segmentation mask helpers, detection/segmentation folder datasets, ResNet-style blocks, MobileNet-style depthwise blocks, a compact UNet, and prediction/mask/feature-map visualization helpers.

import numpy as np
import tensorstudio as ts
from tensorstudio import nn, optim

transform = ts.vision.Compose(
    [
        ts.vision.Resize((8, 8)),
        ts.vision.ToTensor(),
        ts.vision.Normalize(0.5, 0.5),
    ]
)
image = np.zeros((8, 8, 3), dtype=np.uint8)
x = transform(image).reshape((1, 3, 8, 8))

model = ts.vision.ImageClassifier((3, 8, 8), num_classes=2, channels=(4,))
target = ts.tensor([1], dtype="int64")
optimizer = optim.SGD(model.parameters(), lr=0.01)

optimizer.zero_grad()
loss = nn.CrossEntropyLoss()(model(x), target)
loss.backward()
optimizer.step()
print(ts.vision.accuracy(model(x), target))

Image-folder manifests are useful for reproducible dataset snapshots:

manifest = ts.vision.build_image_manifest("data/images", "data/images.manifest.json")
print(manifest["sample_count"])
print(ts.vision.validate_image_manifest("data/images.manifest.json")["valid"])

dataset = ts.vision.ImageManifestDataset("data/images.manifest.json")

Data And Metrics

import tensorstudio as ts
from tensorstudio.data import DataLoader, from_arrays, train_val_split

dataset = from_arrays([[0.0], [1.0], [2.0], [3.0]], [[1.0], [3.0], [5.0], [7.0]])
train_data, val_data = train_val_split(dataset, val_fraction=0.25, seed=42)
loader = DataLoader(train_data, batch_size=2, shuffle=True, seed=42)

for features, targets in loader:
    print(features.shape, targets.shape)

prediction = ts.tensor([[0.2, 0.8], [0.7, 0.3]])
target = ts.tensor([1, 0], dtype="int64")
print(ts.metrics.accuracy(prediction, target))

The v1 DataLoader is intentionally single-process so it works cleanly on Windows without multiprocessing setup.

Projects And Training

tensorstudio.project provides project folders, JSON/TOML/YAML config loading, deterministic seeding, reusable trainers, validation loops, callbacks, safe NPZ weight files, trusted full checkpoints, and generated starter templates:

import tensorstudio as ts
from tensorstudio import nn, optim
from tensorstudio.data import DataLoader, from_arrays, train_val_split
from tensorstudio.project import (
    CSVLogger,
    CheckpointCallback,
    LrLogger,
    Project,
    ProjectConfig,
    Trainer,
    save_state_dict,
    seed_everything,
)

seed_everything(7)
dataset = from_arrays([[0.0], [1.0], [2.0], [3.0]], [[1.0], [3.0], [5.0], [7.0]])
train_data, val_data = train_val_split(dataset, val_fraction=0.25, seed=7)

model = nn.Linear(1, 1)
optimizer = optim.SGD(model.parameters(), lr=0.05)
trainer = Trainer(model, optimizer, nn.MSELoss(), metric_fn=ts.metrics.mean_squared_error)
project = Project("runs/linear", ProjectConfig(name="linear-regression", seed=7))

history = trainer.fit(
    DataLoader(train_data, batch_size=2),
    epochs=50,
    validation_loader=DataLoader(val_data, batch_size=1),
    callbacks=[
        LrLogger(),
        CSVLogger(project.logs_dir / "history.csv"),
        CheckpointCallback(project.checkpoints_dir / "best.tsmodel", save_best_only=True),
    ],
)
save_state_dict(model, project.checkpoint_path("weights"))
print(history.last)

Hardware And Devices

TensorStudio 2.1.0 exposes explicit device descriptors and backend metadata. The published wheels execute tensors on CPU only; CUDA and Metal descriptors are available for feature checks and clear errors.

import tensorstudio as ts

print(ts.available_devices())
print(ts.backend_info())
print(ts.backend_runtime_info())
print(ts.backend_execution_plan("add", "cpu", "float32", input_devices=["cpu", "cpu"]))
print(ts.storage_telemetry())

x = ts.ones((2, 2), device="cpu")
print(x.to("float64").dtype)
print(x.to_device("cpu").device)
print(ts.cuda_is_available())

Passing device="cuda" or device="metal" on CPU-only wheels raises DeviceError instead of silently falling back.

External backend manifests are descriptor-only in this release. TensorStudio can load and validate kernel metadata without importing code or loading shared libraries:

for manifest in ts.discover_kernel_manifests("plugins"):
    ts.register_kernel_manifest(manifest)

Graph Runtime

TensorStudio 1.15.0 adds a constrained graph runtime for supported tensor programs. It traces functions written against symbolic GraphTensor inputs, serializes the graph to JSON, applies simple inspectable optimization passes, and executes the resulting plan through TensorStudio eager tensor operations. It does not trace arbitrary Python control flow and is not a machine-code JIT.

import tensorstudio as ts


def model(x):
    return ((x * 2.0) + 1.0).relu().mean()


graph = ts.trace(model, [ts.TensorSpec((4,), dtype="float32", name="x")])
compiled = ts.compile_graph(graph)
x = ts.tensor([-2.0, -0.25, 1.0, 3.0])

print(compiled(x).item())
print(compiled.profile(x)["nodes"])
print(compiled.memory_plan())

ts.save_graph(graph, "model.tsgraph.json")
loaded = ts.load_graph("model.tsgraph.json")
print(loaded.run(x).item())

Ecosystem Utilities

TensorStudio 2.1.0 expands the late-roadmap ecosystem layer without pretending to be a production-scale distributed or accelerator runtime.

import tensorstudio as ts
from tensorstudio import nn

sparse = ts.sparse_coo_tensor([[0, 1], [1, 0]], [2.0, 3.0], (2, 2))
print(sparse.to_dense().tolist())
print((sparse @ ts.ones((2, 1))).tolist())

model = ts.create_model("tiny_mlp", input_dim=2, hidden_dim=4, output_dim=1)
print(ts.model_info("tiny_mlp")["task"], model(ts.ones((1, 2))).shape)

vocab = nn.Vocabulary.build(["small tensor models", "small language batch"])
inputs, targets = nn.make_causal_lm_batch(vocab.encode("small tensor models"), 2)
lm = nn.CausalLanguageModel(vocab_size=len(vocab), embedding_dim=4, max_length=4)
print(nn.causal_language_model_loss(lm(inputs), targets).item())

quantized = ts.quantization.quantize_tensor(ts.tensor([-1.0, 0.0, 1.0]))
print(quantized.dequantize().tolist())

ts.register_kernel("double", lambda x: x * 2.0, overwrite=True)
print(ts.call_kernel("double", ts.ones((2,))).tolist())
ts.unregister_kernel("double")

for manifest in ts.discover_kernel_manifests("plugins"):
    print(manifest.backend, [kernel.op for kernel in manifest.kernels])

print(ts.distributed.data_parallel_plan(dataset_size=10, batch_size=4))

manifest = ts.data.build_dataset_manifest("data")
print(manifest.validate())

csr = ts.csr_from_dense(ts.tensor([[0.0, 2.0], [3.0, 0.0]]))
print((csr @ ts.ones((2, 1))).tolist())

attention = nn.MultiHeadSelfAttention(embed_dim=4, num_heads=2)
print(attention(ts.randn((1, 3, 4), seed=7), causal=True).shape)

stats = ts.quantization.calibrate_tensor(ts.tensor([-1.0, 0.0, 2.0]))
print(stats.to_dict())

For ONNX files, TensorStudio has two paths:

• ts.import_onnx() imports and executes a constrained static subset through TensorStudio tensor ops.
• ts.run_onnx() can delegate to the optional onnxruntime package when installed with tensorstudio[onnxruntime]; otherwise it can fall back to the supported TensorStudio importer for compatible graphs.
• ts.run_onnx_inference() runs ONNX Runtime with named inputs, named outputs, provider validation, and TensorStudio or NumPy outputs.
• ts.check_onnxruntime_compatibility() and ts.onnxruntime_available_providers() report runtime/provider availability.

outputs = ts.run_onnx_inference(
    "model.onnx",
    {"features": ts.ones((4, 2))},
    output_names="score",
)
print(outputs["score"].shape)

Performance

TensorStudio is optimized for small-to-medium CPU eager workloads, but performance is still experimental. Benchmarks live in benchmarks/ and can be run locally:

python benchmark_all.py
python benchmarks/benchmark_report.py

benchmark_all.py writes benchmarks/results.md and includes explicit win columns for NumPy, TensorFlow, PyTorch, and JAX when those libraries are available locally.

Useful runtime diagnostics:

import tensorstudio as ts

print(ts.performance_info())
ts.set_num_threads(4)

Run the loose local regression thresholds with:

python benchmark_all.py --check-thresholds

On one Windows CPython 3.10 development run reporting 2.1.0, with TensorStudio threads enabled, storage pooling enabled, SSE2 autovectorization reported, and no BLAS provider found, TensorStudio beat NumPy on 8 local benchmark cases and lost on 95 NumPy-comparable cases. JAX CPU dispatch was available on that machine; TensorStudio won 39 local cases and lost 59. The strongest local wins were the simple NumPy convolution/pooling reference loops and some small JAX-dispatch-heavy eager cases. NumPy and JAX were faster for many elementwise, reduction, matrix multiplication, larger activation, and autograd workloads. See benchmarks/results.md for the full table, platform details, and exact timings.

Snapshot from that local run:

operation	shape	TensorStudio	NumPy	JAX CPU dispatch	TS vs NumPy	TS vs JAX
sigmoid	(32,)	0.0316 ms	0.0106 ms	0.1639 ms	0.3360x	5.1782x
mean	(32,)	0.0343 ms	0.0251 ms	0.0261 ms	0.7336x	0.7619x
matmul	(256, 256)	3.4714 ms	0.7755 ms	0.3376 ms	0.2234x	0.0972x
conv2d_3x3_padding1	(1, 1, 8, 8)	0.3263 ms	3.5200 ms	0.1872 ms	10.7889x	0.5738x
avg_pool2d_2x2	(1, 1, 16, 16)	0.0520 ms	1.4566 ms	n/a	28.0110x	n/a
elementwise_backward	(1024,)	5.6504 ms	n/a	n/a	n/a	n/a

Speedup is competitor median / TensorStudio median, so values above 1.0x favor TensorStudio.

Do not treat these results as universal. TensorStudio does not claim to be faster than NumPy, TensorFlow, PyTorch, or JAX overall.

Save And Load

import tensorstudio as ts
from tensorstudio import nn

model = nn.Linear(2, 1)
ts.save({"model": model.state_dict()}, "checkpoint.tsmodel")
checkpoint = ts.load("checkpoint.tsmodel")

Serialization uses pickle. Loading pickle files from untrusted sources is unsafe because pickle can execute arbitrary code.

For safer tensor and state_dict interchange, use TensorStudio's non-pickle NPZ helpers:

state = model.state_dict()
ts.save_npz(state, "weights.tsnpz")
model.load_state_dict(ts.load_npz("weights.tsnpz"))

For safe tensor-only weight files, install the optional SafeTensors extra:

python -m pip install "tensorstudio[safetensors]"

ts.save_safetensors(model.state_dict(), "weights.safetensors")
state = ts.load_safetensors("weights.safetensors")
metadata = ts.inspect_model_metadata("weights.safetensors")

Supported ONNX graphs can be inspected and imported for TensorStudio execution:

ts.export_onnx(model, "model.onnx", input_shape=(1, 2))
print(ts.inspect_onnx("model.onnx")["operators"])
imported = ts.import_onnx("model.onnx")
print(imported(ts.ones((1, 2))).shape)

ONNX Export

TensorStudio can export a supported nn.Sequential graph to ONNX when the optional onnx extra is installed:

import tensorstudio as ts
from tensorstudio import nn

model = nn.Sequential(
    nn.Conv2d(1, 2, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(2 * 2 * 2, 3),
)

ts.export_onnx(model, "classifier.onnx", input_shape=(1, 1, 4, 4))

The exporter supports Linear, Conv2d, grouped/depthwise Conv2d, ConvTranspose2d, Flatten, ReLU, Sigmoid, Tanh, MaxPool2d, and AvgPool2d. TensorStudio also includes ONNX metadata inspection and a supported-subset importer for static graphs. It is not a general ONNX runtime.

Development

python -m pip install -e ".[dev,docs]"
python test_all.py --skip-build
ruff check .
mypy python/tensorstudio
pytest -q
python -m build
python -m twine check dist/*

The native extension module is tensorstudio._C, built with CMake, pybind11, scikit-build-core, and C++20.

Release Checklist

• python test_all.py passes locally.
• ruff check . passes.
• mypy python/tensorstudio passes.
• pytest -q passes on Windows, Linux, and macOS.
• python -m build passes.
• python -m twine check dist/* passes.
• Benchmarks are generated and performance claims match the data.
• Clean wheel installs pass on Windows, Linux, and macOS.
• Clean sdist installs pass on Windows, Linux, and macOS.
• Examples run on all platforms.
• Docs match the implemented feature set.
• No PyPI tokens are committed or printed.
• TestPyPI is verified before a real PyPI release.

Publishing

GitHub Actions build wheels with cibuildwheel. The publish workflow is designed for PyPI trusted publishing with id-token: write; it should not hardcode PyPI tokens or print secrets.

Current Limitations

• CPU backend only.
• Eager execution only.
• No CUDA or Metal tensor execution yet. CUDA/Metal device descriptors and build metadata exist, but unavailable accelerators raise DeviceError.
• Optional BLAS-backed matrix multiplication depends on the build environment exposing a compatible CBLAS/Accelerate interface; otherwise TensorStudio uses a portable C++ fallback.
• No machine-code graph compiler or production distributed runtime. TensorStudio includes a constrained eager-backed graph runtime and single-process distributed planning helpers.
• Convolution and pooling support are CPU-only. Native kernels include NCHW conv2d, grouped/depthwise convolution, conv_transpose2d, max_pool2d, avg_pool2d, and embedding lookup; they are not CUDA/cuDNN replacements.
• Vision covers local image-classification utilities, metrics, visualization, detection/segmentation helpers, compact CNNs, and a compact UNet. It is not an OpenCV replacement and does not include pretrained large model zoos, end-to-end detection/segmentation trainers, video IO, or GPU image kernels yet.
• ONNX support covers export, metadata inspection, import/execution for a limited static subset, and optional delegation to the external ONNX Runtime package through tensorstudio[onnxruntime]. TensorStudio's native importer is not a full ONNX runtime.
• Reductions support all-element, single-axis, and tuple/list-axis reductions for sum, mean, max, and min.
• Arg reductions support all-element flat indices or one axis at a time for argmax and argmin.
• Selection helpers where, maximum, and minimum are native C++ tensor ops with broadcasting and autograd support for floating-point branches.
• Basic integer/slice indexing is supported as native C++ views with autograd scatter-back. Advanced list, tensor, and boolean-mask indexing are not implemented yet.
• Dtype casting is basic and does not include a full promotion/casting policy.
• Experimental performance; benchmarks are local references only.
• Pickle serialization is for trusted TensorStudio objects only.

Roadmap

• CUDA and Metal execution backends
• Graph/JIT mode
• Broader convolution ops, adaptive/global pooling, and image-model examples
• Richer dataset utilities
• Larger model zoo examples and pretrained-weight metadata
• Broader ONNX operator coverage and optional runtime providers
• Runtime-dispatched SIMD kernels
• Better non-BLAS matrix multiplication tiling
• More threaded backward kernels

License

TensorStudio is licensed under the MIT License.