Use 'tensors on device' wording in JAX starters

CLAUDE.md

@@ -145,7 +145,7 @@ Each starter file must have exactly one comment describing the parameters, place
 | CUDA | `// <params> are device pointers` |
 | Mojo | `# <params> are device pointers` |
 | PyTorch, Triton, CuTe | `# <params> are tensors on the GPU` |
-| JAX | `# <params> are tensors on GPU` (+ `# return output tensor directly` inside body) |
+| JAX | `# <params> are tensors on device` (+ `# return output tensor directly` inside body) |
 
 **Rules:**
 - Easy challenges: include the parenthetical `(i.e. pointers to memory on the GPU)` for CUDA/Mojo (matches vector_add reference)
@@ -200,7 +200,7 @@ Verify every item before submitting. This is the single source of truth — work
 - [ ] All 6 files present: `.cu`, `.pytorch.py`, `.triton.py`, `.jax.py`, `.cute.py`, `.mojo`
 - [ ] Exactly 1 parameter description comment per file, no other comments
 - [ ] CUDA/Mojo use "device pointers"; easy challenges include `(i.e. pointers to memory on the GPU)`, medium/hard omit it
-- [ ] Python frameworks use "tensors on the GPU"; JAX also has `# return output tensor directly`
+- [ ] PyTorch/Triton/CuTe use "tensors on the GPU"; JAX uses "tensors on device" and also has `# return output tensor directly`
 - [ ] Starters compile/run but do NOT produce correct output
 
 ### General

challenges/easy/19_reverse_array/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/1_vector_add/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A, B are tensors on GPU
+# A, B are tensors on device
 @jax.jit
 def solve(A: jax.Array, B: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/21_relu/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/23_leaky_relu/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/24_rainbow_table/starter/starter.jax.py

@@ -16,7 +16,7 @@ def fnv1a_hash(x: jax.Array) -> jax.Array:
     return hash_val
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 def solve(input: jax.Array, N: int, R: int) -> jax.Array:
     # return output tensor directly
     pass

challenges/easy/2_matrix_multiplication/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A, B are tensors on GPU
+# A, B are tensors on device
 @jax.jit
 def solve(A: jax.Array, B: jax.Array, M: int, N: int, K: int) -> jax.Array:
     # return output tensor directly

challenges/easy/31_matrix_copy/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A is a tensor on the GPU
+# A is a tensor on device
 @jax.jit
 def solve(A: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/3_matrix_transpose/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, rows: int, cols: int) -> jax.Array:
     # return output tensor directly

challenges/easy/41_simple_inference/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input and model are on the GPU
+# input and model are on device
 def solve(input: jax.Array, model) -> jax.Array:
     # return output tensor directly
     pass

challenges/easy/52_silu/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/54_swiglu/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/62_value_clipping/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, lo: float, hi: float, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/63_interleave/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A, B are tensors on the GPU
+# A, B are tensors on device
 @jax.jit
 def solve(A: jax.Array, B: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/65_geglu/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/66_rgb_to_grayscale/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, width: int, height: int) -> jax.Array:
     # return output tensor directly

challenges/easy/68_sigmoid/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# X is a tensor on GPU
+# X is a tensor on device
 @jax.jit
 def solve(X: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/7_color_inversion/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# image is a tensor on the GPU
+# image is a tensor on device
 @jax.jit
 def solve(image: jax.Array, width: int, height: int) -> jax.Array:
     # return output tensor directly

challenges/easy/8_matrix_addition/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A, B are tensors on the GPU
+# A, B are tensors on device
 @jax.jit
 def solve(A: jax.Array, B: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/easy/9_1d_convolution/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input, kernel are tensors on the GPU
+# input, kernel are tensors on device
 @jax.jit
 def solve(input: jax.Array, kernel: jax.Array, input_size: int, kernel_size: int) -> jax.Array:
     # return output tensor directly

challenges/hard/12_multi_head_attention/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# Q, K, V are tensors on the GPU
+# Q, K, V are tensors on device
 @jax.jit
 def solve(Q: jax.Array, K: jax.Array, V: jax.Array, N: int, d_model: int, h: int) -> jax.Array:
     # return output tensor directly

challenges/hard/14_multi_agent_sim/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# agents is a tensor on the GPU
+# agents is a tensor on device
 @jax.jit
 def solve(agents: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/hard/15_sorting/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# data is a tensor on the GPU
+# data is a tensor on device
 @jax.jit
 def solve(data: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/hard/20_kmeans_clustering/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# data_x, data_y, initial_centroid_x, initial_centroid_y are tensors on the GPU
+# data_x, data_y, initial_centroid_x, initial_centroid_y are tensors on device
 @jax.jit
 def solve(
     data_x: jax.Array,

challenges/hard/36_radix_sort/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/hard/39_Fast_Fourier_transform/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# signal is a tensor on GPU
+# signal is a tensor on device
 @jax.jit
 def solve(signal: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/hard/46_bfs_shortest_path/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# grid is a tensor on the GPU
+# grid is a tensor on device
 @jax.jit
 def solve(
     grid: jax.Array,

challenges/hard/53_casual_attention/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# Q, K, V are tensors on the GPU
+# Q, K, V are tensors on device
 @jax.jit
 def solve(Q: jax.Array, K: jax.Array, V: jax.Array, M: int, d: int) -> jax.Array:
     # return output tensor directly

challenges/hard/56_linear_attention/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# Q, K, V are tensors on the GPU
+# Q, K, V are tensors on device
 @jax.jit
 def solve(Q: jax.Array, K: jax.Array, V: jax.Array, M: int, d: int) -> jax.Array:
     # return output tensor directly

challenges/hard/59_sliding_window_attn/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# Q, K, V are tensors on the GPU
+# Q, K, V are tensors on device
 @jax.jit
 def solve(Q: jax.Array, K: jax.Array, V: jax.Array, M: int, d: int, window_size: int) -> jax.Array:
     # return output tensor directly

challenges/hard/73_all_pairs_shortest_paths/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# dist is a tensor on the GPU
+# dist is a tensor on device
 @jax.jit
 def solve(dist: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/hard/74_gpt2_block/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# x, weights are tensors on GPU
+# x, weights are tensors on device
 @jax.jit
 def solve(x: jax.Array, weights: jax.Array, seq_len: int) -> jax.Array:
     # return output tensor directly

challenges/hard/93_llama_transformer_block/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# x, weights, cos, sin are tensors on GPU
+# x, weights, cos, sin are tensors on device
 @jax.jit
 def solve(
     x: jax.Array,

challenges/medium/10_2d_convolution/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input, kernel are tensors on the GPU
+# input, kernel are tensors on device
 @jax.jit
 def solve(
     input: jax.Array,

challenges/medium/11_3d_convolution/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input, kernel are tensors on the GPU
+# input, kernel are tensors on device
 @jax.jit
 def solve(
     input: jax.Array,

challenges/medium/13_histogramming/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int, num_bins: int) -> jax.Array:
     # return output tensor directly

challenges/medium/16_prefix_sum/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input is a tensor on the GPU
+# input is a tensor on device
 @jax.jit
 def solve(input: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/medium/17_dot_product/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A, B are tensors on the GPU
+# A, B are tensors on device
 @jax.jit
 def solve(A: jax.Array, B: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/medium/18_sparse_matrix_vector_multiplication/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A, x are tensors on the GPU
+# A, x are tensors on device
 @jax.jit
 def solve(A: jax.Array, x: jax.Array, M: int, N: int, nnz: int) -> jax.Array:
     # return output tensor directly

challenges/medium/22_gemm/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# A, B are tensors on the GPU
+# A, B are tensors on device
 @jax.jit
 def solve(
     A: jax.Array, B: jax.Array, M: int, N: int, K: int, alpha: float, beta: float

challenges/medium/25_categorical_cross_entropy_loss/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# logits, true_labels are tensors on the GPU
+# logits, true_labels are tensors on device
 @jax.jit
 def solve(logits: jax.Array, true_labels: jax.Array, N: int, C: int) -> jax.Array:
     # return output tensor directly

challenges/medium/27_mean_squared_error/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# predictions, targets are tensors on the GPU
+# predictions, targets are tensors on device
 @jax.jit
 def solve(predictions: jax.Array, targets: jax.Array, N: int) -> jax.Array:
     # return output tensor directly

challenges/medium/28_gaussian_blur/starter/starter.jax.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 
 
-# input, kernel are tensors on the GPU
+# input, kernel are tensors on device
 @jax.jit
 def solve(
     input: jax.Array,