← All tracks

JAX Profiling & Debugging

Profiler hooks, debug primitives, fixing common tracer/concretization errors, AOT compilation. Practical-skill problems for production work.

0 / 12 solved Continue →
  1. 1. ○ jax.named_scope for Profiler Grouping
  2. 2. ○ profiler.StepTraceAnnotation
  3. 3. ○ @profiler.annotate_function
  4. 4. ○ jax.debug.print with Multiple kwargs
  5. 5. ○ jax.debug.callback for Side Effects
  6. 6. ○ jnp.allclose for Tolerance Equality
  7. 7. ○ Fix: List Append → jnp.concatenate
  8. 8. ○ Fix: Python if → jnp.where
  9. 9. ○ Fix: Python for → vmap
  10. 10. ○ jaxpr with Multiple Args
  11. 11. ○ AOT via jit.lower().compile()
  12. 12. ○ Pretty-Printed jaxpr Length