Skip to content

About

A simple deep learning framework designed for educational purposes and small-scale experiments.

Topics

Resources

Code of conduct

Contributing

Security policy

Stars

6 stars

Watchers

0 watching

Forks

Latest commit

 

History

3,456 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

TensorPlay

Read the whole stack.

A learner-first, DIY-ready framework for tensors, kernels, and custom hardware.

Tutorials • Docs • Quick Start • Installation

EN · 中文

TensorPlay version Service status Python versions License Weekly downloads pip install tensorplay

Build matrix & community signals

Lint English 中文

Platform CPU NVIDIA CUDA AMD ROCm Vulkan Apple MPS

Build Tests

GitHub Stars Last Commit Discord Ask DeepWiki Documentation


TensorPlay is a learner-first deep learning framework. It combines NumPy-style tensor computation with GPU acceleration and neural networks trained by an explicit computation-graph autograd engine. The entire stack — Python API, C++ core, CUDA kernels, compiler — is written to be read end to end: no black box between your model and the hardware.

About TensorPlay

A Transparent Tensor Library

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

Component Description
tensorplay The Python API: tensors, autograd, nn, optim, data loading, serialization — the public surface
p10 C++ core engine: tensor storage and memory management, foundational CPU, CUDA, ROCm/HIP, MPS and Vulkan kernels
tpx Autograd layer: explicit computation-graph construction and backward execution, fully decoupled from the core
stax JIT compiler playground: static graph capture and operator fusion experiments, including native lowering of custom ops
tensorplay.nn Neural network building blocks: Module, Linear, Conv2d, activations, losses, containers — plus the modern LLM stack: rotary embeddings, grouped-query / latent / compressed-sparse attention, delta-rule linear attention, KV caches, autoregressive generation
tensorplay.snn Native spiking neural networks: fused CPU/CUDA LIF/IF neurons, 15 surrogate gradients, multi-step BPTT, spike encoders and deep-SNN layers
tensorplay.quantum Quantum simulation: statevectors and density matrices, gate-table circuits and variational ansätze, measurements and sampling, noise channels, Lindblad master equations and parameter-shift gradients
tensorplay.optim Optimizers (SGD, Adam, AdamW) with learning-rate scheduling and weight decay
tensorplay.utils.data Dataset / DataLoader with multi-worker batching, prefetching and shuffling
tensorplay.library First-class custom operators: register ops, attach fake/meta and autograd formulas, bring your own Triton kernels

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 — 37 dispatch keys"]
    ag --> disp
    disp --> cpu["CPU kernels"]
    disp --> cuda["CUDA kernels"]
    disp --> mps["MPS kernels"]
    disp --> vulkan["Vulkan kernels"]
    disp --> stax["Stax — graph capture & fusion"]
Loading

Supported Backends and Devices

TensorPlay dispatches to five backend families, each with a transparent, readable kernel table:

Backend Target hardware How to enable Check availability
CPU x86_64 / aarch64 CPUs with SIMD dispatch (AVX2/AVX512/NEON) and optional BLAS, oneDNN, MKL, NNPACK and OpenMP acceleration always built in tp.cpu.is_available()
CUDA NVIDIA GPUs — CUDA 12.x / 13.x, cuBLAS/cuDNN USE_CUDA=ON (auto-detected) tp.cuda.is_available()
ROCm / HIP AMD GPUs — ROCm 7.2, MIOpen USE_ROCM=ON (mutually exclusive with CUDA) build-time flag
Vulkan Any Vulkan 1.0+ GPU — discrete NVIDIA/AMD or integrated Intel/AMD iGPU USE_VULKAN=ON (auto-detected) tp.is_vulkan_available()
MPS Apple GPUs on macOS via Metal Performance Shaders — Apple silicon integrated GPU and supported Mac discrete GPUs USE_MPS=ON (auto-enabled on Apple) tp.mps.is_available()

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, CUDA, Vulkan and MPS 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. On macOS the MPS backend is enabled automatically and follows the same readable structure with one Objective-C++/Metal kernel file per operator family.

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.

Spiking Neural Networks (SNN)

tensorplay.snn provides performance-first, native spiking neural network support:

  • Fused neuron kernels — LIF / IF / PLIF / ALIF single-step and multi-step operators are fused CPU/CUDA kernels. Multi-step BPTT runs in a single kernel launch per device; the membrane trajectory is recomputed in backward, so no full forward state is stored.
  • 15 native surrogate gradients — straight-through, sigmoid, fast-sigmoid, piecewise-linear, ATan, Erf, piecewise-quadratic, piecewise-exp, SuperSpike, nonzero-sign-log-abs, piecewise-leaky-ReLU, leaky-K-ReLU, S2NN, q-pseudo-spike and Fourier-series surrogates, evaluated inside the neuron kernels (ids 0-9 in the fused LIF kernels, 0-14 through surrogate_grad).
  • Research-grade controls — per-neuron v_reset broadcasting with gradients, learnable threshold and decay, differentiable v_seq voltage trajectories, soft/hard reset, detach_reset.
  • Deep SNN building blocks — SpikingConv2d, time-folded SpikingBatchNorm1d/2d, SeqToANNContainer, SpikeLinear.
  • Encoders & monitoring — Poisson / latency / weighted-phase / periodic / Gaussian-tuning-curve encoders, SpikeMonitor, spike-rate losses.
  • Performance — the fused multi-step operator reaches ~31 Gspikes/s on CUDA and ~250 Mspikes/s on CPU (114x over a Python time loop).
import tensorplay as tp
from tensorplay import snn

lif = snn.LIFNode(threshold=1.0, decay=0.5, step_mode="m")
out = lif(tp.rand(8, 4, 16))   # [8, 4, 16] spike train with fused BPTT

Quantum Simulation

tensorplay.quantum is a pure-Python quantum simulator for research-sized registers, with the same transparent, autograd-compatible tensors underneath:

  • Circuits & ansätze — a gate-table Circuit covering standard and multi-controlled gates, plus qft_circuit, real_amplitudes, efficient_su2 and nlocal_circuit variational templates with trainable angles.
  • States & data loading — zeros_state / plus_state / basis_state / random_state, encode_state amplitude encoding of classical features, product_state, density_mixture eigendecomposition, partial_trace.
  • Measurements — probabilities, shot sampling, Pauli expectation values, qubit-wise-commuting grouped joint estimation, readout-error correction, and gen_bitstrings basis labels.
  • Noise — single- and multi-qubit Kraus channels (bit_flip, phase_flip, amplitude_damping, depolarizing, arbitrary pauli_kraus models) applied as local operators — no channel superoperator is ever built, and noise rates can be learned parameters.
  • Open-system dynamics — Lindbladian, mesolve, sesolve, trotter_evolve, matrix_exp_evolve, and a standalone dissipator for reading decay rates and steady states.
  • Gradients — param_shift for variational circuits; everything from gate angles to noise rates flows through the autograd engine.
import tensorplay as tp
from tensorplay import quantum as tq

circuit = tq.nlocal_circuit(4, ["ry", "rz"], reps=2)
state = circuit(tq.zeros_state(4))                       # parameterized ansatz
hamiltonian = tq.Hamiltonian([(1.0, "ZZII"), (0.5, "XIZI")])
energy = tq.energy(state, hamiltonian)
energy.backward()                                        # one VQE step

Installation

Binaries

# CPU wheels from PyPI (macOS wheels ship the Apple MPS backend)
pip install tensorplay --upgrade

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

# ROCm wheels for AMD GPUs (ROCm 7.2 / HIP, multi-arch: gfx90a/gfx942/gfx950/gfx1100)
pip install tensorplay \
  --index-url https://download.tensorplay.cn/whl/rocm72-multiarch/ \
  --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. The full variant list lives in .github/cuda-variants.json.

Backends by channel. The macOS CPU wheels carry the MPS (Metal Performance Shaders) backend enabled by default — no separate wheel is needed. ROCm wheels (AMD/HIP) are built as a default release lane and published on their own ROCm pip index (.github/rocm-variants.json).

Tip

The CUDA 13.0 wheels target compute capability 8.6 GPUs only. If your GPU has a different compute capability (or the wheel fails to run on your device), build from source instead and set CMAKE_CUDA_ARCHITECTURES to your target — it defaults to native, which auto-detects the local GPU.

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 — the nightly channel currently builds the cu130 variant
pip install --pre tensorplay \
  --index-url https://download.tensorplay.cn/whl/nightly/cu130/ \
  --extra-index-url https://pypi.org/simple

# CPU nightly — Linux x86_64/aarch64, macOS arm64 (with MPS), Windows x86_64
pip install --pre tensorplay \
  --index-url https://download.tensorplay.cn/whl/nightly/cpu/ \
  --extra-index-url https://pypi.org/simple

# ROCm nightly — AMD GPU lane (rocm72-multiarch)
pip install --pre tensorplay \
  --index-url https://download.tensorplay.cn/whl/nightly/rocm72-multiarch/ \
  --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
  • Apple MPS (Metal Performance Shaders) — enabled automatically on macOS; requires a Metal-capable GPU (Apple silicon or a supported Mac discrete GPU)
  • Ninja (installed automatically with the build dependencies)

Get the TensorPlay Source

git clone https://github.com/lexing-2026/TensorPlay.git
cd TensorPlay
# if you are updating an existing checkout
git submodule sync
git submodule update --init --recursive

The third-party sources ride the submodules, pinned to specific revisions. A checkout that skips them still builds, but the missing pieces degrade: no BLAS vector math, no distributed transports, and the Vulkan backend falls back to CPU unless a system install provides what it needs.

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. A quick check:

python -c "import tensorplay as tp; print(tp.__version__); print('vulkan:', tp.is_vulkan_available()); print('mps:', tp.mps.is_available())"

Tip

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

Run the Test Suite

pytest test/ -q          # everything
pytest test/test_vulkan.py -q   # a single backend suite, e.g. Vulkan
pytest test/test_mps.py -q      # the Apple MPS (Metal) backend suite — macOS only

Adjusting Build Options (Optional)

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

Examples and the full variable table
# CPU-only build
USE_CUDA=OFF pip install .

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

# Apple MPS (Metal) backend — enabled by default on macOS
USE_MPS=ON pip install .

# Target specific GPU architectures
CMAKE_CUDA_ARCHITECTURES="70;75;86" pip install .
Variable Default Description
USE_CUDA auto-detect Enable/disable the CUDA build
USE_ROCM OFF Enable the AMD GPU / HIP build (mutually exclusive with USE_CUDA)
USE_VULKAN auto-detect Enable the Vulkan GPU build; needs a Vulkan loader on the machine plus the vendored headers and GLSL compiler from the submodules
BUILD_TESTS OFF Build the C++ test suite
USE_BLAS / USE_ONEDNN ON BLAS acceleration / oneDNN primitives
MAX_JOBS machine default Cap compile parallelism
DEBUG unset Build with -O0 -g
CMAKE_CUDA_ARCHITECTURES native Semicolon-separated GPU architectures
TENSORPLAY_BUILD_VERSION / TENSORPLAY_BUILD_NUMBER unset Override the package version (used by release CI)

All USE_*, BUILD_* and CMAKE_* environment variables are forwarded to CMake automatically — no extra flags required.

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()

# x.grad is [[11., 15.], [11., 15.]] — the row sums of y, since dL/dz is all ones
print(x.grad)

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 Language Model

tensorplay.nn ships the building blocks modern decoder-only LLMs are made of — rotary embeddings, grouped-query attention, SwiGLU feed-forward, RMSNorm — as composable modules. A TransformerBlock wraps any token mixer (attention, linear attention, …) in the pre-norm residual wiring, so the mixer is the only thing you choose per layer:

import tensorplay as tp
from tensorplay import nn
from tensorplay.nn.cache import DynamicCache


class DecoderLayer(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int, num_kv_heads: int, layer_idx: int):
        super().__init__()
        mixer = nn.GroupedQueryAttention(
            hidden_size, num_heads, num_kv_heads=num_kv_heads,
            qk_norm=True, layer_idx=layer_idx,
        )
        self.block = nn.TransformerBlock(hidden_size, mixer)

    def forward(self, x, position_ids=None, past_key_values=None):
        return self.block(x, position_ids=position_ids, past_key_values=past_key_values)


class TinyLM(nn.Module):
    def __init__(self, vocab_size=256, hidden_size=128, num_layers=4, num_heads=8, num_kv_heads=2):
        super().__init__()
        self.vocab_size = vocab_size
        self.embed = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList(
            [DecoderLayer(hidden_size, num_heads, num_kv_heads, i) for i in range(num_layers)]
        )
        self.norm = nn.RMSNorm(hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, input_ids, position_ids=None, past_key_values=None):
        x = self.embed(input_ids)
        for layer in self.layers:
            x = layer(x, position_ids=position_ids, past_key_values=past_key_values)
        return self.lm_head(self.norm(x))

    def create_empty_cache(self):
        return DynamicCache()


model = TinyLM()
print(model)  # auto-generated architecture visualization

Rotary embeddings and per-head query/key RMS normalization live inside each mixer, and each layer files its key/value pairs into its own slot of the shared cache — the model itself never touches cache internals.

Training Loop

Training a causal LM is ordinary next-token prediction: shift the inputs by one position and apply cross-entropy over the vocabulary.

from tensorplay.nn import CrossEntropyLoss
from tensorplay.optim import AdamW

criterion = CrossEntropyLoss()
optimizer = AdamW(model.parameters(), lr=3e-4)

for step in range(200):
    tokens = (tp.arange(65).unsqueeze(0) + 3 * step) % model.vocab_size  # toy corpus
    logits = model(tokens[:, :-1])          # predict the next token at every position
    loss = criterion(logits.reshape(-1, model.vocab_size), tokens[:, 1:].reshape(-1))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 50 == 0:
        print(f"step {step}: loss {loss.item():.3f}")

Generating Text

generate drives the decode loop: one prefill pass over the prompt, then single-token steps against the cache. Sampling controls — temperature, top-k, nucleus (top-p), repetition penalty, early stop on EOS — are keyword arguments.

from tensorplay.nn.generation import generate

prompt = (tp.arange(8) % model.vocab_size).unsqueeze(0)   # [batch, seq]
out = generate(
    model, prompt,
    max_new_tokens=16, do_sample=True, temperature=0.8, top_k=50,
)
print(out[0].tolist())   # the prompt, continued

Because GroupedQueryAttention shares two KV heads across eight query heads, the cache holds a quarter of the key/value pairs a full multi-head model would. The cache itself is swappable: DynamicCache grows token by token, while StaticCache, SlidingWindowCache, PagedCache, QuantizedCache and the linear-attention LinearStateCache cover static allocation, windowed context, paged serving, quantized storage and recurrent state. Pick one in create_empty_cache() or pass it to generate explicitly.

Beyond GQA, you can swap in MultiheadLatentAttention (low-rank latent KV compression), GatedDeltaNet / KimiDeltaAttention (delta-rule linear attention holding a fixed-size recurrent state per head) or CompressedSparseAttention (sliding window plus compressed KV entries selected by a lightning indexer); CausalEncoderDecoder stacks a decoder over an encoder for encoder–decoder models. For multi-process launching, tprun starts --nproc-per-node copies of your training script wired for distributed runs.

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.

Interpretability with lumen

tensorplay.lumen is the framework-native explainability suite: classic attribution (integrated gradients, DeepLIFT, LRP, SHAP, TCAV), layer and neuron attribution, and a set of research-grade LLM tools — sparse autoencoders, cross-layer transcoders, attribution graphs, feature circuits, logit lens, and causal tracing. Everything runs on tensorplay primitives; no third-party runtime is required.

from tensorplay.lumen import IntegratedGradients
from tensorplay.lumen.llm import ActivationPatching, SparseAutoencoder

# feature attribution on any nn.Module
attr = IntegratedGradients(model).attribute(prompt, target=1, n_steps=50)

# causal tracing: which layer and position stores the fact?
trace = ActivationPatching(model).causal_trace(clean_ids, corrupted_ids, target_position=-1)

A step-by-step guide reproducing methods from the interpretability literature (causal tracing, logit lens, SAE features, cross-layer transcoders, attribution graphs, feature circuits) is in the lumen guide.

Metrics

tensorplay.metrics is a ready-to-use evaluation suite covering the common machine-learning domains, split into stateful Metric classes and stateless functional functions:

  • Classification: accuracy, precision/recall/F1, ROC/AUC, precision-recall curve, Matthews correlation coefficient, Cohen's kappa, calibration error, hinge, Jaccard, specificity, LogAUC, threshold-based recall/precision curves
  • Regression: MSE/RMSE, MAE, R², Pearson/Spearman/Kendall correlations, KL/JS divergence, concordance correlation, Tweedie deviance, critical success index, and more
  • Image: PSNR, SSIM, total variation, and AIGC metrics (FID, KID, Inception Score, LPIPS)
  • Audio: SNR/SI-SNR, SDR/SI-SDR/SA-SDR with permutation-invariant training, PESQ, STOI, DNSMOS, NISQA, SRMR
  • Text: BLEU, ROUGE, WER/CER, CHRF, TER, SacreBLEU, BERTScore, SQuAD, InfoLM, perplexity
  • Retrieval: AUROC, MAP, MRR, NDCG, precision/recall, R-precision, fall-out, hit-rate, recall-at-fixed-precision
  • Nominal: Cramér's V, Fleiss' kappa, Pearson's contingency coefficient, Theil's U, Tschuprow's T
  • Remote sensing: ERGAS, RASE, SAM, UIQI, sliding-window RMSE, spectral distortion indices
  • Wrappers: MetricCollection, MetricTracker, running metrics, classwise/multioutput/min-max wrappers, bootstrapping and multitask wrappers

Metrics mirror the reference API, so existing workflows port with minimal changes: instantiate a metric, call update(preds, target) repeatedly, then compute(), or use the functional form directly. Metrics that need optional third-party packages (transformers, pesq, pystoi, onnxruntime, ...) are importable out of the box and raise a clear ModuleNotFoundError until the dependency is installed.

A fused dual-input reduction (sum_squared_error / sum_abs_error) powers the regression metrics: one vectorized pass computes sum((a-b)^2) / sum(|a-b|) without materialising intermediate tensors, making MSE and MAE several times faster than the reference implementation on CPU (see Benchmarks).

Benchmarks

TensorPlay benchmarks: per-op throughput spread, backend op coverage, dispatch overhead

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, MPS operator latencies
  • Reports: every script emits a JSON report (--json-out) — throughput, latency percentiles, compile cost — ready for plotting and for the white paper evaluation

Metric kernels use fused reductions where possible. On CPU (AVX2, 8 threads), selected regression metrics beat the reference implementation by a wide margin:

Metric TensorPlay Reference Speedup
MSE (10k×64) 0.01 ms 0.04 ms 3.3×
MAE (10k×64) 0.01 ms 0.04 ms 2.8×
Kendall tau-b (n=200) 4.8 ms 124 ms 26×

Measured highlights from the white paper, §9 — a snapshot of current main: the dispatcher adds the same sub-1% sliver on CPU, CUDA and Vulkan paths; the CUDA backend covers 1,262 unique ops (97% of the CPU surface, 1,233 ops); the Vulkan teaching backend ships 224 ops backed by 6.8k lines of GLSL shaders across 102 files. The MPS backend adds Metal-accelerated kernels on macOS — tensor ops, convolutions, attention/LLM kernels (RMSNorm, SDPA, top-k sampling, fused AdamW) — with its own operator benchmark and a CI performance floor (bench_mps_operator.py, mps_benchmark.yml). The ROCm lane builds and smoke-tests AMD/HIP wheels in the ROCm container on every release and nightly run. The pytest suite defines 3,660 test functions.

Scripts and methodology: benchmark/README.md.

Testing

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

pytest test/

The metrics suite adds about 184 tests across test/test_metrics*.py — stateful metrics, functional forms, curves and plots, AIGC, audio, retrieval, nominal, remote sensing, wrappers and lazy-loaded external-dependency metrics.

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; ROCm on Linux x86_64) and validates every wheel on every pull request and push to main. macOS arm64 wheels ship the MPS backend, with dedicated Metal acceptance and benchmark workflows; the ROCm lane runs in the ROCm build container on the same release/nightly pipelines. See .github/workflows/ (pull, trunk, publish, nightly, lint). Lint rules live under [tool.ruff] in pyproject.toml.

Resources

Communication

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 — CPU (macOS wheels with the MPS backend), CUDA (cu124/cu126/cu130) and ROCm (AMD/HIP) lanes — publishes CPU wheels to PyPI, and uploads the CUDA and ROCm wheels to a GitHub Release with a PEP 503 index on Cloudflare Pages. Wheel-matrix infrastructure (platform runners, CUDA variants, the 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.

Contributors

License

TensorPlay is licensed under the Apache 2.0 License.

Star History Chart

The Organization: TensorPlay

ohtensorplay 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

Dependency
NumPy
SymPy
TensorBoard
Pillow
pybind11

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

About

A simple deep learning framework designed for educational purposes and small-scale experiments.

Topics

Resources

Code of conduct

Contributing

Security policy

Stars

6 stars

Watchers

0 watching

Forks

Releases

Packages

Used by

Contributors

Languages