TensorPlay is a Python package that provides two high-level features:

• Tensor computation (like NumPy) with strong GPU acceleration
• Deep neural networks built on an explicit, tape-based autograd system

The whole stack — Python API, C++ core, CUDA kernels, compiler — is engineered to be read: clean implementations, explicit computation graphs, and no black box between your model and the hardware.

• About TensorPlay
  • A Transparent Tensor Library
  • Why TensorPlay
• Installation
  • Binaries
  • From Source
    • Prerequisites
    • Get the TensorPlay Source
    • Install Build Dependencies
    • Install TensorPlay
    • Adjusting Build Options (Optional)
• Getting Started
  • Automatic Differentiation
  • Defining a Neural Network
  • Training Loop
• Benchmarks
• Testing
• Resources
• Communication
• Releases and Contributing
• License
• The organization behind TensorPlay

About TensorPlay

A Transparent Tensor Library

At a granular level, TensorPlay consists of the following components:

The four pillars — P10, TPX, Stax and NN — are deliberately decoupled libraries that can work together or independently. Domain subpackages such as linalg, fft, sparse, special, amp, distributed and serialization round out the API surface.

Every call is one short, visible path — no black box between your model and the hardware:

flowchart TB
    py["Python API — tensorplay / nn / optim / data"] --> ag["TPX autograd — explicit DAG"]
    py --> disp["P10 dispatcher — 13 dispatch keys"]
    ag --> disp
    disp --> cpu["CPU kernels"]
    disp --> cuda["CUDA kernels"]
    disp --> vulkan["Vulkan kernels"]
    disp --> stax["Stax — graph capture & fusion"]

Why TensorPlay

TensorPlay is built on a philosophy of transparency: you can trace every operation from Python into the C++ core without getting lost in abstraction layers.

• Pure and readable implementations. Dive into the logic of every operator — from autograd to memory management. No black boxes.
• DIY acceleration. Simplified CPU and CUDA backends are a playground for experimenting with custom hardware kernels and learning parallel computing.
• Modular autograd. The decoupled TPX engine builds computation graphs explicitly, making backpropagation easy to understand and extend.
• Research ready. Prototype new layer types, optimizers, custom operators and storage formats with minimal boilerplate.
• Instantly familiar API. If you have worked with mainstream deep learning frameworks, the mental model carries over — spend your time on internals, not syntax.

Familiar and Python-First

TensorPlay is not a thin binding over an opaque engine. The public surface — tensorplay, nn, optim, data — is pure Python you can read end to end, and every call crosses a small, explicit boundary into the C++ core. Stack traces, error messages and the debugger show your code, not framework plumbing.

DIY Hardware Acceleration

The CPU and CUDA backends are deliberately simplified: one kernel family per file, one registration macro per unit, no hidden scheduling. That makes them a playground for learning parallel computing and prototyping kernels — bring your own device backend and wire it into the dispatcher without touching the framework itself.

An Explicit Autograd Engine

The autograd engine (TPX) is a standalone library, fully decoupled from the core. Computation graphs are built and executed explicitly — no hidden state — so backpropagation is something you can read, trace and extend in an afternoon.

Extensions Without Pain

First-class custom operators: register an op with tensorplay.library, attach fake/meta formulas for shape inference and an autograd formula for backward, and it composes with the whole stack — including native lowering into the Stax compiler or your own Triton kernels.

Installation

Binaries

# CPU wheels from PyPI
pip install tensorplay --upgrade

# CUDA wheels from the TensorPlay CUDA index (choose cu124, cu126, or cu130)
# Keep PyPI as an extra index for runtime dependencies.
pip install tensorplay \
  --index-url https://download.tensorplay.cn/whl/cu124/ \
  --extra-index-url https://pypi.org/simple

[!NOTE] Make sure your Python version matches the wheel tags (e.g. cp310 for Python 3.10). For CUDA wheels, the driver and runtime must support the CUDA version of the wheel.

Nightly (preview) builds

Try tomorrow's features today: every change that passes our build-and-smoke pipeline lands on the rolling nightly channel automatically, following the nightly version format (X.Y.0.dev<date>+cuXXX / +cpu). Only the latest build per variant is kept.

# CUDA nightly (choose cu124, cu126, or cu130)
pip install --pre tensorplay \
  --index-url https://download.tensorplay.cn/whl/nightly/cu124/ \
  --extra-index-url https://pypi.org/simple

# CPU nightly
pip install --pre tensorplay \
  --index-url https://download.tensorplay.cn/whl/nightly/cpu/ \
  --extra-index-url https://pypi.org/simple

From Source

Building from source gives you a hackable, debuggable install — the recommended setup for contributors and anyone working on kernels.

Prerequisites

• Python >= 3.10, < 3.14
• CMake >= 3.18 (< 4.0)
• A C++20-capable compiler (MSVC 2022 on Windows, GCC/Clang on Linux)
• CUDA Toolkit (optional, for GPU support); set CMAKE_CUDA_ARCHITECTURES to target specific GPU architectures
• ROCm 7.2.x (optional, for AMD GPU support); the HIP backend is currently built from source
• Ninja (installed automatically with the build dependencies)

Get the TensorPlay Source

git clone https://github.com/lexing-2026/TensorPlay.git
cd TensorPlay

Install Build Dependencies

Isolated PEP 517 builds fetch everything automatically. For faster iterative development, install the toolchain once and build without isolation:

pip install -r requirements-build.txt

Install TensorPlay

TensorPlay is built with scikit-build-core through the standard PEP 517 interface declared in pyproject.toml:

# Full build and install (add -v for verbose output)
pip install .

# Editable install for development
pip install -e . --no-build-isolation

# Or produce a wheel without installing
python -m build --wheel

On success, import tensorplay picks up the compiled _C extension from the installed package.

[!TIP] For day-to-day kernel work, the editable install (pip install -e . --no-build-isolation) recompiles only what changed.

Adjusting Build Options (Optional)

Environment variables drive the build — no -D flags needed:

# CPU-only build
USE_CUDA=OFF pip install .

# AMD GPU build with ROCm 7.2 / HIP
USE_CUDA=OFF USE_ROCM=ON pip install .

# Target specific GPU architectures
CMAKE_CUDA_ARCHITECTURES="70;75;86" pip install .

Getting Started

Automatic Differentiation

import tensorplay as tp

x = tp.Tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
y = tp.Tensor([[5.0, 6.0], [7.0, 8.0]], requires_grad=True)

z = x.matmul(y) + tp.ones_like(x)
loss = z.sum()
loss.backward()

print(x.grad)  # [[6., 6.], [6., 6.]]

Under the hood, TPX records each operation into an explicit DAG and replays the chain rule node by node: $\dfrac{\partial \mathcal{L}}{\partial x} = \dfrac{\partial \mathcal{L}}{\partial z} \cdot \dfrac{\partial z}{\partial x}$ — every edge of that graph is code you can step through.

Defining a Neural Network

import tensorplay as tp
from tensorplay.nn import Module, Linear, ReLU, Sigmoid

class MLP(Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        self.fc1 = Linear(input_dim, hidden_dim)
        self.relu = ReLU()
        self.fc2 = Linear(hidden_dim, output_dim)
        self.sigmoid = Sigmoid()

    def forward(self, x: tp.Tensor) -> tp.Tensor:
        x = self.relu(self.fc1(x))
        return self.sigmoid(self.fc2(x))

model = MLP(10, 32, 1)
print(model)  # auto-generated architecture visualization

Training Loop

from tensorplay.data import DataLoader, TensorDataset

train_data = TensorDataset(tp.randn(100, 10), tp.randn(100, 1))
train_loader = DataLoader(dataset=train_data, batch_size=8, shuffle=True)

for batch_x, batch_y in train_loader:
    predictions = model(batch_x)
    # ... compute loss, call loss.backward(), step the optimizer

Structured tutorials — linear regression from scratch, MNIST CNN classification, custom datasets, model saving/loading with .mega + state_dict — live at tensorplay.cn. Deep dives, one pillar per post — dispatch, autograd engine, tensor storage, compiler — live in the blog series.

Benchmarks

The benchmark/ suite measures what a readable framework costs — and proves correctness at the same time. Every comparison runs both runtimes from identical initial weights, deterministic data order, and the same optimizer and batch settings; the gate is strict: logits must match within allclose tolerance and Top-1 predictions must be identical before any timing is reported.

• End-to-end training: ResNet-18 image classification on CPU and CUDA, eager and compiled (stax, native lowering), reporting train/eval/test accuracy and throughput
• Micro and subsystem: GEMM, optimizer steps, dataloader, serialization, autograd Function overhead, custom-op call overhead, LLaMA end-to-end
• Reports: every script emits a JSON report (--json-out) — throughput, latency percentiles, compile cost — ready for plotting and for the white paper evaluation

Measured highlights (see the white paper, §9): the dispatcher adds the same sub-1% sliver on CPU, CUDA and Vulkan paths; the CUDA backend covers 1,274 unique ops (96% of the CPU surface); the Vulkan teaching backend ships 145 ops backed by 4.5k lines of GLSL shaders.

Scripts and methodology: benchmark/README.md.

Testing

The Python test suite runs with pytest against an installed (or built in-place) package:

pytest test/

CI builds the full wheel matrix (Python 3.10–3.13; CPU on Linux x86_64/aarch64, macOS arm64, and Windows x86_64; CUDA on Linux and Windows x86_64) and validates every wheel on every pull request and push to main; see .github/workflows/ (pull, trunk, publish, lint). Lint rules live under [tool.ruff] in pyproject.toml.

Resources

• Documentation: tensorplay.cn
• Tutorials: tensorplay.cn/en/guide/tutorials
• White Paper: docs/whitepaper/main.pdf — the whole stack with path:line citations
• Blog: docs/blogs — one pillar per post
• Benchmarks: benchmark/
• Community: Discord

Communication

• Discord: discord.gg/u6T5e2kGJm — questions, ideas, showcase
• GitHub Issues: bug reports, feature requests, RFCs
• Email: feedback@tensorplay.cn

Releases and Contributing

The package version lives in version.txt (single source of truth; development installs carry a +git<sha> suffix). Releases are cut by pushing a vMAJOR.MINOR.0 tag; patch tags such as v1.2.1 do not publish binaries. CI then builds and validates the full wheel matrix, publishes CPU wheels to PyPI, and uploads CUDA wheels to a GitHub Release with a PEP 503 index on Cloudflare Pages. Wheel-matrix infrastructure (platform runners, CUDA variants, the optional ROCm channel) is documented in RELEASE.md.

We welcome contributions of all kinds — bug fixes, documentation, new features. See CONTRIBUTING.md for the development workflow and coding standards.

License

TensorPlay is licensed under the Apache 2.0 License.

The Organization: TensorPlay

ohtensorplay Open tools and infrastructure for people who want to understand, experiment with, and build AI systems. Organization mission · Make Everyone a Great AI-Builder

Projects

• megatensors — the MEGA model-hub SDK powering .mega serialization, the mega:// storage backend, and hub model loading (pip install megatensors)

Other dependencies

Explore all ohtensorplay projects — give it a ⭐ and follow along.