[pests]

I will say the

  jax.lax.max()

call in the example is unintentionally funny to me. Almost felt like I was reading code written by Dr. Seuss.