Selam Waktola

Hi @Lime-Cakes It looks like this issue appears to be resolved. The Kaggle notebook you provided (https://www.kaggle.com/code/liminchen1/update-jax-tpu-bug) executed successfully and printed local_devices in approximately 7.1287e-05 seconds. Could you please confirm...

Hi @GeorgOstrovski To illustrate the use of `device_put_sharded` within JIT-compiled code, I created an example. When executed with JAX version 0.4.23 on both CPU and GPU, ```Python import jax import...

> Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll...

Hi @jakevdp Is the following modification a doable solution for inserting embedded code examples (customized to JAX) for NumPy and SciPy? We would use `lax_description` for inserting examples and leave...

It make sene, docstring is rewritten from scratch and tried to squash the commits.

Thanks @jakevdp 'See also' section added and indentations corrected. Locally, the formatting appears as follows: 

Thanks @jakevdp - Description is adjusted: - `ndarray` is modified to "Tuple of JAX arrays". - integer sequences is modified to integer arrays. - New example is added.

unnecessary texts and example outputs are modified/removed.

> Looks like there are some lint issues with whitespace. You can fix these locally by running `pre-commit`: https://jax.readthedocs.io/en/latest/contributing.html#linting-and-type-checking It seems the lint failure was below: `trim trailing whitespace.................................................Failed` removing...

Hi @SkalskiP, Based on this conversation at [X](https://twitter.com/__kolesnikov__/status/1791073793403727981), you can use [numpy.savez](https://numpy.org/doc/stable/reference/generated/numpy.savez.html#numpy.savez), a function that efficiently stores multiple arrays in a single, uncompressed `.npz` format. ```Python import big_vision.utils as bv_utils...