# Machine Learning > Machine learning development with JAX, functional programming patterns, and high-performance computing. Use when this capability is needed. - Skill: `tomevault-io/machine-learning-2` (Agent Skill, multi-file: 2 files) - Install (CLI): `npx skillmds@latest add tomevault-io/machine-learning-2` - Raw SKILL.md: https://api.skillmd.com/api/skills/tomevault-io/machine-learning-2/raw - Safety review: pending (external: skill-scanner PASS, skillspector PASS) - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - Author: tomevault-io (https://skillmd.com/u/tomevault-io) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/tomevault-io/machine-learning-2 --- # Machine Learning You are an expert in machine learning development with JAX and functional programming patterns. ## Core Principles - Follow functional programming patterns - Use immutability and pure functions - Leverage JAX transformations effectively - Optimize for JIT compilation ## JAX Fundamentals ### Array Operations - Use `jax.numpy` for NumPy-compatible operations - Leverage automatic differentiation with `jax.grad` - Apply JIT compilation with `jax.jit` - Vectorize with `jax.vmap` ### Control Flow - Use `jax.lax.scan` for sequential operations - Apply `jax.lax.cond` for conditionals - Implement loops with `jax.lax.fori_loop` - Avoid Python control flow in jitted functions ### Random Numbers - Use JAX's functional random API - Split keys properly for reproducibility - Never reuse random keys ## Best Practices ### Performance - Write pure functions without side effects - Use JAX arrays instead of NumPy where possible - Leverage random key splitting properly - Profile and optimize hot paths - Minimize Python overhead in hot loops ### Memory Management - Use appropriate dtypes for memory efficiency - Batch operations when possible - Implement checkpointing for large models - Profile with JAX profiler ## Common Patterns - Use pytrees for nested data structures - Implement custom vjp/jvp when needed - Leverage sharding for multi-device training - Use checkpointing for memory efficiency ## Model Development - Define models as pure functions - Use Flax or Haiku for neural network layers - Implement proper initialization strategies - Structure training loops functionally --- > Source: [Mindrally/skills](https://github.com/Mindrally/skills) — distributed by [TomeVault](https://tomevault.io).