Jake Vanderplas comments

Results 637 comments of


                                            Jake Vanderplas

Vmap of dynamic_update_slice very slow on TPU

Hi, thanks for the report! I think this performance discrepancy comes from the fact that the batching rule for `dynamic_slice` is implemented in terms of `gather`, and `gather` is much...

Explicit managing the output buffer of `jax.jit` function

You may be able to do what you want via the `donate_argnums`/`donate_argnames` parameter of `jax.jit`; see the [`jax.jit` documentation](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) for a description. For example: ```python from functools import partial import...

Explicit managing the output buffer of `jax.jit` function

There is a new experimental `ArrayRef` object that may serve the purposes of the question here: https://docs.jax.dev/en/latest/array_refs.html This should be part of the v0.7.1 release.

Explicit managing the output buffer of `jax.jit` function

Closing becuase I don't think there's any other action to take here. Thanks for raising the issue!

Performance Discrepancy in `jax.value_and_grad` for Different `argnums`

It's hard to say much definitive without a full reproduction, including what operations `state.apply_fun` are doing, but in general it's not surprising that taking derivatives with respect to different inputs...

Performance Discrepancy in `jax.value_and_grad` for Different `argnums`

I'm having trouble running your code: can you point to what the `modeling` package is? These details help, but my first response still holds: differentiating with respect to different arguments...

docs: `default_device()` can be `None`

A comment here: the wording implies that `None` has a special meaning – but my understanding is that the spec makes no requirements of what objects libraries use to represent...