neuron-explorer
Steering-first neuron explorer for Qwen3.5-4B: push every MLP neuron, keep every rollout, browse the results
TensorPlay is a transparent deep learning framework for learning, research, and hardware experiments. Trace tensors from Python through autograd, dispatch, and custom kernels.
Developer toolsPython · C++GitHub @lexing-2026OnlineOwner unconfirmed
Have you tried it? One tap is enough. No sign-in needed.
Owners of public repositories under a personal GitHub account can sign in to verify ownership and manage the introduction themselves. No administrator approval is needed.
Sign in, generate a verification file, and commit it to this repository’s default branch within 30 minutes. Files on other branches or forks do not count. Then check the file below to manage the project without administrator approval. Keep the file to retain your access.
If needed, you will go to sign-in. After signing in, return here and press your chosen verification button again.
Run /nomorevibe verify in your project folder to verify ownership and manage the introduction yourself.
Verify in one minuteAsk about other verification options
A simple deep learning framework designed for educational purposes and small-scale experiments.
Taken from the repository README
Loading the full README. The saved excerpt is shown below.
Read the whole stack. A learner-first, DIY-ready framework for tensors, kernels, and custom hardware. EN · 中文 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. At a granular level, TensorPlay consists of the following components:Read the full README ↗
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.
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"]
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() |
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.
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.
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.
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.
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.
tensorplay.snn provides performance-first, native spiking neural network support:
surrogate_grad).v_reset broadcasting with gradients, learnable threshold and decay, differentiable v_seq voltage trajectories, soft/hard reset, detach_reset.SpikingConv2d, time-folded SpikingBatchNorm1d/2d, SeqToANNContainer, SpikeLinear.SpikeMonitor, spike-rate losses.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
tensorplay.quantum is a pure-Python quantum simulator for research-sized registers, with the same transparent, autograd-compatible tensors underneath:
Circuit covering standard and multi-controlled gates, plus qft_circuit, real_amplitudes, efficient_su2 and nlocal_circuit variational templates with trainable angles.zeros_state / plus_state / basis_state / random_state, encode_state amplitude encoding of classical features, product_state, density_mixture eigendecomposition, partial_trace.gen_bitstrings basis labels.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.Lindbladian, mesolve, sesolve, trotter_evolve, matrix_exp_evolve, and a standalone dissipator for reading decay rates and steady states.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
# 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.
cp310for 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_ARCHITECTURESto your target — it defaults tonative, which auto-detects the local GPU.
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
Building from source gives you a hackable, debuggable install — the recommended setup for contributors and anyone working on kernels.
CMAKE_CUDA_ARCHITECTURES to target specific GPU architecturesgit 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.
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
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.
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
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 .
# 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.
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.
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 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}")
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.
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.
tensorplay.metrics is a ready-to-use evaluation suite covering the common machine-learning domains, split into stateful Metric classes and stateless functional functions:
MetricCollection, MetricTracker, running metrics, classwise/multioutput/min-max wrappers, bootstrapping and multitask wrappersMetrics 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).
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.
stax, native lowering), reporting train/eval/test accuracy and throughput--json-out) — throughput, latency percentiles, compile cost — ready for plotting and for the white paper evaluationMetric 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.
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.
path:line citationsThe 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.
TensorPlay is licensed under the Apache 2.0 License.
.mega serialization, the mega:// storage backend, and hub model loading (pip install megatensors)| Dependency |
|---|
| NumPy |
| SymPy |
| TensorBoard |
| Pillow |
| pybind11 |
Explore all ohtensorplay projects — give it a ⭐ and follow along.