rajasekharporeddy comments

Results 25 comments of


                                            rajasekharporeddy

process can exit before jax.debug.prints, needs jax.effects_barrier()

Hi @jbuckman @mattjj Looks like this issue is resolved in JAX version 0.4.26. I executed the mentioned code on colab (T4 GPU) with JAX version 0.4.26 without adding the `jax.effects_barrier()`...

Print large value with suppress in printoptions

Hi @sh0416 Thanks for the question. JAX leverages NumPy's implementation of `set_printoptions` directly by assigning it to `jax.numpy.set_printoptions`. https://github.com/google/jax/blob/c4dea624cce95e2dcd288831eadd68d1aae4a05a/jax/_src/numpy/lax_numpy.py#L117C1-L117C39 From https://github.com/numpy/numpy/issues/21653#issuecomment-1145751948, it can be understood that numpy displays the numbers...

`jax.lax.with_sharding_constraint` is not propagated outside jit

Hi @Findus23 It looks the issue mentioned by you has been resolved. I tried to reproduce the issue mentioned by you on Colab with JAX version 0.4.23. Now the operations...

Nesting `vmap` within `pmap` using `jax.pure_callback` segfaults

Hi @Justin-Tan Looks like this issue has been resolved in later versions of JAX. I executed the provided repro code with JAX version 0.4.23 on Google Colab using GPU run...

JIT+Shardmap+donate_argnums doesn't work, while pmap+donate_argnums works.

Hi @nouiz IIUC.. Looks like this issue has been resolved. I tried to reproduce the issue on Google colab (both CPU and GPU) with JAX version 0.4.23. The code executed...

`roll` incorrect if jit-compiled on GPU with 64-bit mode

Hi @mathisgerdes Looks like this issue has been resolved in later versions of JAX. I executed the mentioned code on colab (GPU T4) with cuda 12.3 and cuDNN 8.9.7 and...

`roll` incorrect if jit-compiled on GPU with 64-bit mode

Hi @mathisgerdes Please feel free to close the issue, if it is resolved. Thank you.

Add the ability to pass an `is_leaf` argument to `tree_reduce`.

Hi @gileshd It looks like this issue has been resolved with PR #16018. I tried to execute the mentioned using JAX version 0.4.23 on colab by importing `tree_reduce` and `tree_map`...

all_to_all with axis_index_group argument broken

Hi @mganahl It looks like this has been resolved. I executed the mentioned code with JAX version 0.4.23 on Colab CPU and JAX version 0.3.25 on colab TPU. Now the...

JIT+Shardmap+donate_argnums doesn't work, while pmap+donate_argnums works.

Hi @wbrenton I tested it on colab TPU v2 with JAX version 0.4.26 and JAX nightly version 0.4.31.dev20240630 and it works fine. Attaching [gist](https://colab.sandbox.google.com/gist/rajasekharporeddy/01d55a78d45ece94f4b3e6df1034cf08/tpu_15368.ipynb) for reference. Thank you.