Paul Jeha comments

Results 7 comments of


                                            Paul Jeha

Specify non-pytree node dataclass fields

Would love to see this feature brought to Chex :)

Batchnorm example doesn't work

This seems to work now, no ? At least on the colab. Has anything changed since ?

Incompatibility with tensorflow 2.16

I installed tensorflow 2.15, tensorflow_probability 0.23 and then tensorflow_gan and it worked

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

Hi, I was wondering what the status was on that ? I face the following situation: ```python jax.vmap(lambda x,y: lax.cond(y

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

Hey, thanks for the suggestion, I create a MWE that I think illustrates your idea. Let me know if that is not the case. This example batches the condition as...

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

Thanks a lot, I'm a bit surprised by your results. Your `false_fn` is very cheap compared to `true_fn` so I would have expected that vmapping would have been faster than...

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

I've been observing the same on gpu :/