xla-agentic-development

Plugin xla-agentic-development from marketplace alpharomercoma in alpharomercoma/xla-agentic-development. Commands read from the repo's manifest — 1 month ago.

Skills (9)

• tpu-colab-provisioning

  Get a working TPU in about a minute using the official Google Colab CLI, and run code on it headlessly from a local terminal. Covers session lifecycle, running scripts and notebooks non-interactively, probing what the runtime actually ships, and the Kaggle and Cloud TPU spot fallbacks when one chip

• tpu-env-setup

  Provision and verify a Google TPU development environment. Covers Compute Engine and GKE provisioning, jax[tpu] / torch_xla / libtpu installation, PJRT configuration, and non-destructive verification. Use when the user says "set up TPU", "create a TPU VM", "provision v6e", "jax doesn't see my TPU",

• tpu-hlo-artifacts

  Dump and read XLA compiler artifacts for TPU: HLO text, proto and HTML dumps, StableHLO from jax.jit lowering, optimized HLO from compilation, cost_analysis FLOP and byte estimates, and the persistent compilation caches. Use when the user says "dump HLO", "xla_dump_to", "StableHLO", "what did XLA co

• tpu-pallas-debugging

  Diagnose and fix Pallas TPU kernel failures: Mosaic compilation errors, vector-layout and tiling errors, VMEM exhaustion, DMA and semaphore races, hangs, and wrong-number bugs. Use when the user says "Mosaic error", "pallas kernel won't compile", "apply-vector-layout", "not divisible by 128", "VMEM

• tpu-pallas-docs

  Look up Pallas TPU and TPU architecture documentation: API signatures, memory spaces, pipelining primitives, tiling rules, and per-generation hardware specs. Use when the user asks "what does pl.BlockSpec do", "pltpu API", "signature of make_async_copy", "how much VMEM does v5p have", "MXU size on v

• tpu-pallas-writing

  Write and modify Pallas TPU kernels. Covers new kernels from PyTorch/NumPy/JAX or a natural-language description, editing existing kernels, adding shape and dtype support, refactoring grid and BlockSpec tiling, and calling Pallas kernels from PyTorch/XLA. Use when the user says "write a Pallas kerne

• tpu-profile-analysis

  Analyze captured XProf TPU profiles: roofline placement and bound-by classification, HLO op ranking, step-time decomposition, memory peaks, and known anti-pattern detection. Produces a ranked bottleneck diagnosis with evidence. Use when the user says "analyze this trace", "roofline", "HLO op profile

• tpu-profiling

  Capture TPU execution profiles with XProf for JAX and PyTorch/XLA workloads, and enumerate devices with tpu-info. Use when the user says "profile my TPU job", "capture a trace", "xprof", "jax.profiler", "start_trace", "xp.start_trace", "tensorboard profile TPU", "tpu-info", "how busy is my TPU", or

• tpu-torch-xla-lowering

  Diagnose PyTorch/XLA performance and correctness problems on TPU: CPU fallbacks from unlowered ops, graph recompilation from dynamic shapes, and lazy-tensor graph-boundary mistakes. Use when the user says "torch_xla is slow", "aten fallback", "PT_XLA_DEBUG", "metrics_report", "recompiling every step

Subagents (4)

• tpu-pallas-agent

  Unified Pallas TPU kernel agent covering the full lifecycle: writing kernels from PyTorch, NumPy, JAX or a natural-language description; debugging Mosaic compilation errors, layout failures, VMEM exhaustion and wrong numbers; profiling; and looking up Pallas and TPU hardware documentation. Routes to

• tpu-pallas-debugger-agent

  Autonomously diagnoses and fixes failing Pallas TPU kernels: Mosaic compilation errors, vector-layout and tiling failures, VMEM exhaustion, DMA and semaphore races, hangs, and wrong numerical results. Triages into a failure class, reproduces under interpret mode, applies the class-appropriate fix, a

• tpu-pallas-writer-agent

  Authors and modifies Pallas TPU kernels. Translates PyTorch, NumPy, JAX, or a natural-language description into a correct Pallas kernel; adds shape and dtype support; refactors grid and BlockSpec tiling. Verifies against a CPU reference but does not tune performance. <example> user: "Convert this Nu

• tpu-performance-agent

  Diagnoses TPU performance problems end to end: captures XProf profiles, classifies the bottleneck from roofline data, inspects HLO artifacts, and hunts PyTorch/XLA CPU fallbacks and recompilation. Produces a ranked, evidenced diagnosis. Read-only by design — it measures and explains, it does not edi

Hooks (1)

• SessionStart

  "${CLAUDE_PLUGIN_ROOT}"/hooks/scripts/check-tpu-env.sh

Similar plugins