Skip to content
cccskills
FreePublished on GitHub

jax-best-practices

Expert in JAX for high-performance numerical computing and machine learning

The English summary is being prepared. Showing the original description.

See how to install

Files included (1)

  • SKILL.md1.3 KB

SKILL.md (original)

Review the instructions the agent will receive before you install.

JAX Best Practices

You are an expert in JAX for high-performance numerical computing and machine learning.

Core Principles

  • Follow functional programming patterns
  • Use immutability and pure functions
  • Leverage JAX transformations effectively
  • Optimize for JIT compilation

Key Transformations

jax.jit

  • Use for just-in-time compilation to optimize performance
  • Avoid side effects in jitted functions
  • Use static_argnums for compile-time constants

jax.vmap

  • Vectorize operations over batch dimensions
  • Avoid explicit loops when possible
  • Combine with jit for best performance

jax.grad

  • Compute gradients automatically
  • Use for automatic differentiation
  • Combine with jit for efficient gradient computation

Best Practices

  • Write pure functions without side effects
  • Use JAX arrays instead of NumPy where possible
  • Leverage random key splitting properly
  • Profile and optimize hot paths

Performance

  • Minimize Python overhead in hot loops
  • Use appropriate dtypes
  • Batch operations when possible
  • Profile with JAX profiler

Common Patterns

  • Use pytrees for nested data structures
  • Implement custom vjp/jvp when needed
  • Leverage sharding for multi-device
  • Use checkpointing for memory efficiency

Reviews

No reviews yet. Tell others what you thought.

Comments

No login needed. Share how it worked for you, ask a question, or note where it did not run.

No comments yet. Be the first to write one.

Comments are published right away. Do not include personal information such as email addresses or phone numbers. Up to 2 URLs.

0 / 1,000 characters

This form is protected by reCAPTCHA, and the Google Privacy Policy and Terms of Service apply.

More skills from this repository

12-agent academic paper writing pipeline. 10 modes (full/plan/outline/revision/revision-coach/abstract/lit-review/format-convert/citation-check/disclosure). 6 paper types, 5 citation formats, bilingual abstracts, LaTeX/DOCX-via-Pandoc/PDF output. Style Calibration + Writing Quality Check + Anti-Patterns with IRON RULE markers. Triggers: write paper, academic paper, guide my paper, parse reviews, AI disclosure, 寫論文, 學術論文, 引導我寫論文, 審查意見.

The English summary is being prepared. Showing the original description.

bouclem/skills7Updated May 31, 2026

Multi-perspective academic paper review with dynamic reviewer personas. Simulates 5 independent reviewers (EIC + 3 peer reviewers + Devil's Advocate) with field-specific expertise. Supports full review, re-review (verification), quick assessment, methodology focus, Socratic guided, and calibration modes. Triggers on: review paper, peer review, manuscript review, referee report, review my paper, critique paper, simulate review, editorial review, calibrate reviewer, reviewer calibration, measure reviewer accuracy.

The English summary is being prepared. Showing the original description.

bouclem/skills7Updated May 31, 2026

Orchestrator for the full academic research pipeline: research -> write -> integrity check -> review -> revise -> re-review -> re-revise -> final integrity check -> finalize. Coordinates deep-research, academic-paper, and academic-paper-reviewer into a seamless 10-stage workflow with mandatory integrity verification, two-stage peer review, and reproducible quality gates. Triggers on: academic pipeline, research to paper, full paper workflow, paper pipeline, end-to-end paper, research-to-publication, complete paper workflow.

The English summary is being prepared. Showing the original description.

bouclem/skills7Updated May 31, 2026

Create, iterate, and scale paid ad creative for Google Ads, Meta, LinkedIn, TikTok, and similar platforms. Use when generating headlines, descriptions, primary text, or large sets of ad variations for testing and performance optimization.

The English summary is being prepared. Showing the original description.

bouclem/skills7Updated May 31, 2026

This skill should be used when the user asks to "implement LLM-as-judge", "compare model outputs", "create evaluation rubrics", "mitigate evaluation bias", or mentions direct scoring, pairwise comparison, position bias, evaluation pipelines, or automated quality assessment.

The English summary is being prepared. Showing the original description.

bouclem/skills7Updated May 31, 2026

AI agent development workflow for building autonomous agents, multi-agent systems, and agent orchestration with CrewAI, LangGraph, and custom agents.

The English summary is being prepared. Showing the original description.

bouclem/skills7Updated May 31, 2026

See all skills by bouclem

Report a problem with this skill