Roy Frostig

icon indicating a website link https://cs.stanford.edu/~rfrostig/

icon company @google-deepmind icon location sfba icon location ML researcher, JAX co-author

Results 34 comments of Roy Frostig

vmap of cond's predicate results in select, leading to unexpected compute/memory use

We're in agreement here by and large. This is something that we've thought about improving before, whether at the JAX or XLA level. I can't find an open issue for...

vmap of cond's predicate results in select, leading to unexpected compute/memory use

@minqi – the thinking hasn't changed much since this issue was last active. Although there's a fundamental puzzle regarding whether/how to do better, for now we're still producing `select` when...

vmap of cond's predicate results in select, leading to unexpected compute/memory use

> @inversecrime would sequential_vmap be similar to unrolling the batch using a for loop? It generates a `jax.lax.map`, which bottoms out in a (rolled) XLA loop.

Support Hessian of gamma-distributed samples

We believe that we don't need `random_gamma_grad_p` as a primitive at all any more, as proposed in #27628.