jax

206d2f98 - [host_callback] Fix handling of params[donated_invars] (#4040)

[host_callback] Fix handling of params[donated_invars] (#4040) There was some confusion when rewriting xla_call primitives and the donated_invars was not set to be equal to the length of the invars. Fixes: #4015 Fixes: #3863