jax-skills

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.

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.

python

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.

python

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).

python

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.

python

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.

python

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.

python

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.

python

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

python

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.