jax

d7269489 - [Pallas][GPU kernel interpreter] Clarify type annotations, esp. of Jax/Numpy arrays, of callback arguments.