rajasekharporeddy
rajasekharporeddy
Hi @Gattocrucco The fix provided by openxla [PR #10413](https://github.com/openxla/xla/pull/10413) appears to address this issue. I tested the mentioned code with JAX version 0.4.26 in Google Colab. The results produced by...
Hi @FelipeMartins96 This issue appears to be addressed by PR #8195 which removed the unused variable `param_scale` in Neural Network docs. Thank you.
Hi @mattjj Looks like this issue has been resolved in later versions of JAX. I tried to reproduce this issue with the latest JAX version 0.4.26 on cloud VM with...
Hi @SGENZO Thank for reporting the bug. I have opened a PR #20772 on this. This issue will be closed once the PR is merged.
Hi @neel04 I tested the provided reproducible code on my Mac-book Pro with M1 Pro chip using jax versions 0.4.26 and 0.4.27.dev20240503 and the corresponding jaxlib versions 0.4.26 and 0.4.27.dev20240503,...
Hi @bloops I executed the mentioned code with JAX version 0.3.25 on Google Colab TPU. It executed without any error. ```python import jax.numpy as jnp from jax.experimental.maps import xmap from...
Hi @yitongx It looks like `jax.tree_util.tree_reduce` has been added to the JAX documentation with the PR #19588. The same is reflected in the documentation here: https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.tree_reduce.html. Could you please verify...
Hi @marcvanzee I executed the mentioned code on Colab using JAX version 0.4.23. The error message reads as follows: ```python ValueError: pmap got inconsistent sizes for array axes to be...
Hi @mattjj It looks like the error message has been improved by including `f"Found slice({i.start}, {i.stop}, {i.step}) "` in the error message by the PR #8124. Please refer https://stackoverflow.com/questions/76611278/indexerror-array-slice-indices-must-have-static-start-stop-step-to-be-used-with to...
Hi @hawkinsp Using jax-metal 0.0.6 with jax 0.4.26 and jaxlib 0.4.23, the code now encounters an `XLARuntimeError` instead of crashing with an `Abort trap 6` when dealing with the unsupported...