jax

jax is a skill for Claude Code, Codex from G1Joshi/Agent-Skills. It costs 14 tokens per session (312 once invoked), scanned A, original, MIT.

A Python-based numerical-computing system that combines NumPy-style calculations with automatic calculation of derivatives and compilation for faster execution. It is used for machine learning and scientific research.

In plain words
What is it for?
Use it for machine-learning research, high-order derivatives, compiled numerical functions, and training across GPUs or Google TPUs.
Why use it?
It can speed up repeated numerical code and calculate derivatives for complex mathematical or model-training work.

Skill for Claude CodeCodex

Written for no agent in particular: nothing here depends on one.

Good fit Use it for machine-learning research, high-order derivatives, compiled numerical functions, and training across GPUs or Google TPUs.

Compare 6 skills from other repositories ↓
Install with agentmods
npx agentmods add skills/g1joshi/agent-skills/jax
Install

Getting it into your agent

One page per mod, every tool's command on it. A separate URL per tool would split the same page into five that compete with each other.

Any agent
npx skills add G1Joshi/Agent-Skills --skill jax
Clone the repo
git clone --depth 1 https://github.com/G1Joshi/Agent-Skills

Made for: Claude Code, Codex.

Wrote this? Show the measurements

A badge with what this costs and how it scanned, read live from this page, so it follows the numbers instead of freezing them. Markdown for a README, HTML for a documentation site or a project page.

agentmods badge for jax

README.md
[![agentmods](https://agentmods.dev/badge/skills/g1joshi/agent-skills/jax/github.svg)](https://agentmods.dev/skills/g1joshi/agent-skills/jax)
Your own site
<a href="https://agentmods.dev/skills/g1joshi/agent-skills/jax"><img src="https://agentmods.dev/badge/skills/g1joshi/agent-skills/jax/github.svg" alt="Measured on agentmods" height="20"></a>

Or the 80×15 button, for a site that already has a row of RSS and ATOM ones. Only the verdict fits; the numbers stay here.

agentmods 80×15 button for jax

Your own site · 80×15
<a href="https://agentmods.dev/skills/g1joshi/agent-skills/jax"><img src="https://agentmods.dev/badge/skills/g1joshi/agent-skills/jax.svg" alt="Reviewed on agentmods" width="80" height="20"></a>
Per session 14 Skills are progressive disclosure: only the name and description are preloaded; the body loads when the skill is used.
When invoked 312 The whole file, excluding the scripts and references it only reads on demand.
Security scan A 0 findings. A grade says what 26 rules found in the file — not that it is safe.
Origin original No closer match found in the catalogue.
Token cost

What it costs to keep this loaded

Counted locally with the o200k_base tokenizer, which is exact for GPT models; Claude uses its own tokenizer and its counts differ. Treat this as one consistent yardstick across the catalogue rather than a bill. Prices are per million input tokens.

ModelPer sessionOnce invoked
Fable 5.1 $0.00014 $0.00312
Opus 5 $0.00007 $0.00156
Sonnet 5 $0.00003 $0.00062
Haiku 4.5 $0.00001 $0.00031

Measured 9d ago against content hash 410a24148029, method: parsed. Prices are Anthropic first-party input rates as of 2026-09-09, from the pricing page.

Security

Grade A, and why

jax scanned grade A with 0 findings against 26 rules in 11 categories — prompt injection, anti-refusal, data exfiltration, privilege escalation, supply chain, agent snooping, system-prompt leakage, SSRF and excessive agency — measured 9d ago.

A static scan of the body, not an audit. Every finding is printed with the line that produced it so you can judge whether it matters here. A mod is markdown that instructs an agent; that is exactly why what it instructs is worth reading.

Nothing flagged

None of the 26 patterns this scan looks for appear in this file: no shell pipes, no recursive deletes, no credential paths, no hidden text, no instruction-override or anti-refusal phrasing, no agent-config snooping. That is not a guarantee, it is the absence of the things that are checkable.

skills/ai-ml/jax/SKILL.md · 45 lines

What it actually says

JAX

JAX is "NumPy on steroids". It combines Autograd (automatic differentiation) with XLA (compilation). 2025 sees Flax NNX (PyTorch-style OOP) becoming standard.

When to Use

  • TPU Training: JAX runs natively on Google TPUs.
  • Research: If you need to compute 10th order derivatives or strange math.
  • Massive Scale: DeepMind and OpenAI use JAX for training frontier models.

Core Concepts

Functional Transformations

grad(), jit(), vmap(), pmap().

Flax (NNX)

Neural network library. NNX introduces mutable state (OOP) to make JAX feel like PyTorch.

Statelessness

(Legacy Flax) parameters are stored separately from the model.

Best Practices (2025)

Do:

  • Use jit: Always compile your functions.
  • Use Flax NNX: Avoid the complexity of legacy immutable Flax/Haiku.
  • Use shard_map: For distributed training across devices.

Don't:

  • Don't use side effects: print() inside a jit function only runs once (during tracing).

References

Changes

What this file has done since we first saw it

Hashed on every crawl. A supply-chain change to an agent config is a question of when, not whether, so the history is kept rather than the latest state alone.

  1. 9d ago First seen · 45 lines · 14 tokens per session scan A 410a24148029

Subscribe to this mod's changes

jax is a skill published in the GitHub repository G1Joshi/Agent-Skills (12 stars, last pushed 7mo ago), licensed MIT. It adds 14 tokens to every session and 312 once invoked, about $0.0001 per session on Opus 5. A static security scan graded it A with 0 findings. No closer match exists in the catalogue, so it is treated as the original; first seen 2026-08-30.

Related

Other skills, from other repositories