jax-skills

•

9382423 months ago

High-performance numerical computing and machine learning workflows using JAX. Supports array operations, automatic differentiation, JIT compilation, RNN-style scans, map/reduce operations, and gradient computations. Ideal for scientific computing, ML models, and dynamic array transformations.

Quick Install

npxskills add benchflow-ai/skillsbench--skill jax-skills

Instructions

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.

hseq = jx.rnn_scan(sequence, Wx, Wh, b)

Notes:

Returns sequence of hidden states.
Uses tanh activation.

5. JIT Compilation

`jit_run(fn, args)`

Description: JIT compile and run a function using JAX. Parameters:

fn (callable): Function to compile.
args (tuple): Arguments for the function.

result = jx.jit_run(my_function, (arg1, arg2))

Notes:

Speeds up repeated function calls.
Input shapes must be consistent across calls.

Best Practices

Prefer JAX arrays (jnp.array) for all operations; convert to NumPy only when saving.
Avoid side effects inside functions passed to vmap or scan.
Validate input shapes for map_op, reduce_op, and rnn_scan.
Use JIT compilation (jit_run) for compute-heavy functions.
Save arrays using .npy or pickle/json to avoid system-specific issues.

Example Workflow

import jax.numpy as jnp
import jax_skills as jx

# Load array
arr = jx.load("data.npy")

# Square elements
arr2 = jx.map_op(arr, "square")

# Reduce along axis
mean_arr = jx.reduce_op(arr2, "mean", axis=0)

# Compute logistic gradient
grad_w = jx.logistic_grad(X_train, y_train, w_init)

# RNN scan
hseq = jx.rnn_scan(sequence, Wx, Wh, b)

# Save result
jx.save(hseq, "hseq.npy")

Notes

This skill set is designed for scientific computing, ML model prototyping, and dynamic array transformations.
Emphasizes JAX-native operations, automatic differentiation, and JIT compilation.
Avoid unnecessary conversions to NumPy; only convert when interacting with external file formats.

Loading…

Related Skills

skill-creator

Create or update AgentSkills. Use when designing, structuring, or packaging skills with scripts, references, and assets.

367,360

skill-creator

Create or update AgentSkills. Use when designing, structuring, or packaging skills with scripts, references, and assets.

367,309

1password

Set up and use 1Password CLI (op). Use when installing the CLI, enabling desktop app integration, signing in (single or multi-account), or reading/injecting/running secrets via op.

365,915

himalaya

CLI to manage emails via IMAP/SMTP. Use `himalaya` to list, read, write, reply, forward, search, and organize emails from the terminal. Supports multiple accounts and message composition with MML (MIME Meta Language).

365,885

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.

hseq = jx.rnn_scan(sequence, Wx, Wh, b)

Notes:

Returns sequence of hidden states.
Uses tanh activation.

5. JIT Compilation

`jit_run(fn, args)`

Description: JIT compile and run a function using JAX. Parameters:

fn (callable): Function to compile.
args (tuple): Arguments for the function.

result = jx.jit_run(my_function, (arg1, arg2))

Notes:

Speeds up repeated function calls.
Input shapes must be consistent across calls.

Best Practices

Prefer JAX arrays (jnp.array) for all operations; convert to NumPy only when saving.
Avoid side effects inside functions passed to vmap or scan.
Validate input shapes for map_op, reduce_op, and rnn_scan.
Use JIT compilation (jit_run) for compute-heavy functions.
Save arrays using .npy or pickle/json to avoid system-specific issues.

Example Workflow

import jax.numpy as jnp
import jax_skills as jx

# Load array
arr = jx.load("data.npy")

# Square elements
arr2 = jx.map_op(arr, "square")

# Reduce along axis
mean_arr = jx.reduce_op(arr2, "mean", axis=0)

# Compute logistic gradient
grad_w = jx.logistic_grad(X_train, y_train, w_init)

# RNN scan
hseq = jx.rnn_scan(sequence, Wx, Wh, b)

# Save result
jx.save(hseq, "hseq.npy")

Notes

This skill set is designed for scientific computing, ML model prototyping, and dynamic array transformations.
Emphasizes JAX-native operations, automatic differentiation, and JIT compilation.
Avoid unnecessary conversions to NumPy; only convert when interacting with external file formats.