Update docs of callbacks

[roth-jakob]

Callback functions should not call into JAX. This information was missing in the docs of the callbacks. This commit adds this information to the docs.

See: #25861, #24255

[Edenhofer]

Thank you @roth-jakob for hunting down the cause of the hangs in NIFTy with recent JAX versions!! Personally, I think this information is worth a big warning box in the doc-string but I might be alone in this view. I think the innocent looking a == b case with a and b being JAX arrays is also worth mentioning explicitly in the docs. It is obvious that this calls back into JAX but I think it is very easy to overlook. Maybe it is even worth suggesting to users to convert JAX arrays to numpy arrays if they intend to do any math and/or comparisons on arguments to a callback. WDYT?

[stephen-huan]

Unrelated to this PR but on the topic of callback documentation changes, shouldn't jax.debug.callback have guaranteed execution? It doesn't make sense otherwise. (That is, the last column of the jax.debug.callback row should have a ✅ instead of a ❌.)

jax/docs/external-callbacks.md

Lines 76 to 80 in 4f8699c

	|callback function | supports return value | `jit` | `vmap` | `grad` | `scan`/`while_loop` | guaranteed execution |
	|-------------------------------------|----|----|----|----|----|----|
	|{func}`jax.pure_callback` | ✅ | ✅ | ✅ | ❌¹ | ✅ | ❌ |
	|{func}`jax.experimental.io_callback` | ✅ | ✅ | ✅/❌² | ❌ | ✅³ | ✅ |
	|{func}`jax.debug.callback` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ |

I can open another PR if need be, sorry for hijacking this one.

[dfm]

Thanks for this @roth-jakob! I think this is a good place to start, but hopefully we can improve the default behavior soon.

[dfm]

@stephen-huan — You might look at this discussion for a conversation about your question. @sharadmv's comment there says that JAX has permission to eliminate calls to debug.callback in some cases. I must admit that I'm not totally sure what those cases are, but it might be better to revive that thread or open a new issue for this discsusion.

[stephen-huan]

@dfm I see, thanks for answering my question. Perhaps that should be another documentation fix, in another PR.

[stephen-huan]

For what it's worth, I find the wording "must not include any calls back into JAX" a bit confusing, especially when talking about callbacks ("callbacks" vs "calls back"). I think something like "must not involve any computation with JAX arrays" a bit more clear (with some intentional ambiguity with what "computation" means exactly). My mental model is anything that one can access statically (e.g. .shape, .dtype, etc.) is ok but computation that would be staged out when jit'd is dangerous. I haven't actually tested this intuition.