Skip to content
← Learn JAX

JAX Conceptual Deep-Dives

JAX

Custom pytree nodes, host interop callbacks, quantization patterns, jaxpr inspection, advanced debugging. Beyond the curriculum.

This track uses JAX. You can make it your practice mode to browse related content across the site.

15 exercises

Explore freely. Sign in to keep your progress.

Start track

Your learning sequence

  1. 1 Not solved yet. register_pytree_node Practice · Medium
  2. 2 Not solved yet. tree_flatten / tree_unflatten Round-Trip Practice · Medium
  3. 3 Not solved yet. Pytree Aux Data Practice · Hard
  4. 4 Not solved yet. pure_callback Basics Practice · Medium
  5. 5 Not solved yet. io_callback for Side Effects Practice · Hard
  6. 6 Not solved yet. jax.debug.print Practice · Easy
  7. 7 Not solved yet. JAX and NumPy Array Conversion Practice · Easy
  8. 8 Not solved yet. int8 Affine Quantization Practice · Medium
  9. 9 Not solved yet. int8 Dequantization Practice · Easy
  10. 10 Not solved yet. float8 Round-Trip Practice · Medium
  11. 11 Not solved yet. Counting Jaxpr Equations Practice · Medium
  12. 12 Not solved yet. jax.debug.callback Practice · Medium
  13. 13 Not solved yet. Selective stop_gradient Practice · Hard
  14. 14 Not solved yet. JVP with Two Input Directions Practice · Hard
  15. 15 Not solved yet. jaxpr with jit Practice · Medium

Check yourself

4 questions · not saved while you're signed out

Use these questions to check your understanding and compare explanations. Your score does not affect track completion.

0 / 4 correct

What does jax.tree.map do with a nested dict of arrays?

Choose one answer to reveal the explanations. Answers are kept only while this page is open. 

tree = {"a": jnp.array(1.0), "b": [jnp.array(2.0), jnp.array(3.0)]}
jax.tree.map(lambda v: v * 2, tree)
Question 1 of 4