jax-skills

jax-skills is a skill for Claude Code, Codex from benchflow-ai/skillsbench. It costs 54 tokens per session (1,090 once invoked), scanned A, original, Apache-2.0.

A numerical-computing toolkit built around JAX, a Python library for fast array calculations and machine-learning operations. It covers array handling, automatic differentiation, just-in-time compilation, and mapped or repeated computations.

In plain words
What is it for?
Use it to load and save arrays, apply array operations, calculate gradients, compile computations, and build scientific or machine-learning workflows.
Why use it?
It provides consistent ways to run numerical workflows, validate array inputs, and save or transform data in JAX-compatible formats.

Skill for Claude CodeCodex

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

Good fit Use it to load and save arrays, apply array operations, calculate gradients, compile computations, and build scientific or machine-learning workflows.

Compare 6 skills from other repositories ↓
Install with agentmods
npx agentmods add skills/benchflow-ai/skillsbench/jax-skills
About the project

SkillsBench is a benchmark for measuring how effectively AI agents use modular skills—folders containing instructions, scripts, and resources—to complete specialized tasks. It helps researchers and developers evaluate both skill quality and agent behavior, including tasks that require combining multiple skills. The catalogue’s skills and instructions are evaluated as part of this workflow.

benchflow-ai/skillsbench · 1,760 stars · on GitHub · skillsbench.ai

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 benchflow-ai/skillsbench --skill jax-skills
Clone the repo
git clone --depth 1 https://github.com/benchflow-ai/skillsbench

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-skills

README.md
[![agentmods](https://agentmods.dev/badge/skills/benchflow-ai/skillsbench/jax-skills/github.svg)](https://agentmods.dev/skills/benchflow-ai/skillsbench/jax-skills)
Your own site
<a href="https://agentmods.dev/skills/benchflow-ai/skillsbench/jax-skills"><img src="https://agentmods.dev/badge/skills/benchflow-ai/skillsbench/jax-skills/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-skills

Your own site · 80×15
<a href="https://agentmods.dev/skills/benchflow-ai/skillsbench/jax-skills"><img src="https://agentmods.dev/badge/skills/benchflow-ai/skillsbench/jax-skills.svg" alt="Reviewed on agentmods" width="80" height="20"></a>
Per session 54 Skills are progressive disclosure: only the name and description are preloaded; the body loads when the skill is used.
When invoked 1,090 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. Third-party audits
  • NVIDIA SkillSpector pass 7 Sept 2026
How audits are shown
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.00054 $0.01090
Opus 5 $0.00027 $0.00545
Sonnet 5 $0.00011 $0.00218
Haiku 4.5 $0.00005 $0.00109

Measured 7d ago against content hash 2cff1fe8e449, method: parsed. Prices are Anthropic first-party input rates as of 2026-09-11, from the pricing page.

Security

Grade A, and why

jax-skills 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 7d ago.

The scan reads SKILL.md. This mod also ships 1 executable file (jax_skills.py), listed below but not scanned — reading those needs a real analyzer, not pattern matching.

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.

Origin

Copies of this mod

2 near-identical copies found in the catalogue:

tasks/jax-computing-basics/environment/skills/jax-skills/SKILL.md · 153 lines

How it starts

The opening of the file, as written. The whole thing — 153 lines — stays where its author put it; the contents beside it link to each section on GitHub.

Requirements for Outputs

General Guidelines

Arrays

  • All arrays MUST be compatible with JAX (jnp.array) or convertible from Python lists.
  • Use .npy, .npz, JSON, or pickle for saving arrays.

Operations

  • Validate input types and shapes for all functions.
  • Maintain numerical stability for all operations.
  • Provide meaningful error messages for unsupported operations or invalid inputs.

JAX Skills

1. Loading and Saving Arrays

load(path)

Description: Load a JAX-compatible array from a file. Supports .npy and .npz.
Parameters:

  • path (str): Path to the input file.

Returns: JAX array or dict of arrays if .npz.

import jax_skills as jx

arr = jx.load("data.npy")
arr_dict = jx.load("data.npz")

save(data, path)

Description: Save a JAX array or Python array to .npy. Parameters:

  • data (array): Array to save.
  • path (str): File path to save.
jx.save(arr, "output.npy")

2. Map and Reduce Operations

map_op(array, op)

Description: Apply elementwise operations on an array using JAX vmap. Parameters:

  • array (array): Input array.
  • op (str): Operation name ("square" supported).
squared = jx.map_op(arr, "square")

reduce_op(array, op, axis)

Description: Reduce array along a given axis. Parameters:

  • array (array): Input array.
  • op (str): Operation name ("mean" supported).
  • axis (int): Axis along which to reduce.
mean_vals = jx.reduce_op(arr, "mean", axis=0)

3. Gradients and Optimization

logistic_grad(x, y, w)

Description: Compute the gradient of logistic loss with respect to weights. Parameters:

  • x (array): Input features.
  • y (array): Labels.
  • w (array): Weight vector.
grad_w = jx.logistic_grad(X_train, y_train, w_init)

Notes:

  • Uses jax.grad for automatic differentiation.
  • Logistic loss: mean(log(1 + exp(-y * (x @ w)))).

4. Recurrent Scan

rnn_scan(seq, Wx, Wh, b)

Description: Apply an RNN-style scan over a sequence using JAX lax.scan. Parameters:

  • seq (array): Input sequence.
  • Wx (array): Input-to-hidden weight matrix.
  • Wh (array): Hidden-to-hidden weight matrix.
  • b (array): Bias vector.

Read the full file on GitHub · 153 lines

Files

What ships with it

8 files beside SKILL.md in the same directory: the scripts, references and assets a skill reads on demand. Not counted in the per-session cost; read them before you install if any of them is executable.

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. 7d ago First seen · 153 lines · 54 tokens per session scan A 2cff1fe8e449

Subscribe to this mod's changes

jax-skills is a skill published in the GitHub repository benchflow-ai/skillsbench (1,760 stars, last pushed 1mo ago), licensed Apache-2.0. It adds 54 tokens to every session and 1,090 once invoked, about $0.0003 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-09-03.

Related

Other skills, from other repositories

dbscan-custom-metric

Run DBSCAN clustering with a custom distance metric using sklearn, including how to define weighted Euclidean metrics and extract cluster centroids.

cxcscmu/SkillLearnBench · 32 tokens

netcdf-processing

Reading, processing, and analyzing NetCDF output from lake simulation models.

cxcscmu/SkillLearnBench · 17 tokens

Multimodal Alignment

Align speech, text, image, or video signals for multimodal benchmarks.

Raidriar7170/hermes-skilleval · 20 tokens

embodied-eval-automation

Plan, explain, build, run, monitor, validate, transfer, and audit reproducible embodied-model studies and batch episode collection. Use when a user wants to connect a local, SSH, or cloud GPU host; understand and compare a policy, VLA, world model, world-action model, or hybrid with a benchmark; discover and reuse…

Yinzhanqing/embodied-eval-automation · 145 tokens

torchdrug

Build and troubleshoot TorchDrug 0.2.1 workflows for molecular graphs, property prediction, self-supervised pretraining, molecule generation, retrosynthesis, protein representation learning, and knowledge graph reasoning. Use when code imports torchdrug or needs its datasets, models, tasks, or Engine.

K-Dense-AI/scientific-agent-skills · 61 tokens

deepspot-m

Generate transcriptome-wide virtual spatial transcriptomics from H&E histology with DeepSpot-M. Use when you need spatial gene expression in log1p-CPM for 224x224 tiles at about 20x, want to query protein-coding genes by symbol instead of a fixed panel, or want to run prediction across a whole slide after tiling with…

K-Dense-AI/scientific-agent-skills · 80 tokens