evo-jax-computing-basics

evo-jax-computing-basics

End-to-end skill for solving JAX computation tasks: parses problem.json, loads input data, dispatches to the correct JAX computation, and saves results as .npy files.

Quick Start

# Run directly from the command line:
python /app/environment/skills/evo-jax-computing-basics/scripts/solve_tasks.py /app/problem.json

Or from Python:

import sys
sys.path.insert(0, '/app/environment/skills/evo-jax-computing-basics/scripts')
from solve_tasks import solve_all_tasks

solve_all_tasks('/app/problem.json')

Key Functions

I/O

• load_problem_json(path) - Load and parse problem.json, returns list of task dicts
• load_input_data(file_path) - Load .npy or .npz file, convert arrays to JAX arrays
• save_jax_array(path, jax_array) - Safely save JAX array to .npy file using np.asarray()

Compute

• compute_basic_reduce(x) - Row-wise mean of 2D array using jnp.mean(x, axis=1)
• compute_map_square(x) - Element-wise square using jnp.square(x)
• compute_grad_logistic(data) - Gradient of logistic loss w.r.t. weights using jax.grad
• compute_scan_rnn(data) - RNN forward pass using jax.lax.scan, returns stacked hidden states (T, H)
• compute_jit_mlp(data) - JIT-compiled 2-layer MLP forward pass with ReLU activation
• dispatch_task(task_id, data) - Dispatch to appropriate compute function by task ID (with fuzzy matching)

Orchestration

• solve_all_tasks(problem_path) - Full pipeline: load problems, compute, save results

Computation Details

basic_reduce

• Input: 2D array x of shape (M, N)
• Output: 1D array of shape (M,) with mean of each row
• Uses: jnp.mean(x, axis=1)

map_square

• Input: array x of any shape
• Output: array of same shape with each element squared
• Uses: jnp.square(x) (natively vectorized, no need for vmap)

grad_logistic

• Input: dict with x (N,D), y (N,) with labels in {-1,1}, w (D,)
• Loss: L(w) = mean(logaddexp(0, -y * (X @ w)))
• Output: gradient of loss w.r.t. w, shape (D,)
• Uses: jax.grad with argnums=0, jnp.logaddexp(0, z) for numerical stability