Issue 1:
I am seeing this on my local install of jax 0.3.13 with jaxlib 0.3.10

import jax.numpy as jnp
r = jnp.ones(3)
r[10]
DeviceArray(1., dtype=float32)

Shouldn't r[10] return an out-of-bound IndexError instead?

Check for duplicate issues.

Issue 2:
GPU/TPU cannot be found after jax process is imported in another open session

Please:

Check for duplicate issues.

This looks different - running 1 piece of code on multiple devices

Concerned here, unsure if "we haven't published an API for multi-host GPU" is relevant here.

jax version 0.2.17

jaxlib version 0.1.67

Steps to replicate:

On a system with multiple devices (in my case, google cloud v2-8 TPU VM).

Open a terminal.

>>> import jax
>>> jax.devices()
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]

With that python session still open, open another terminal/tmux session in my case.

export TF_CPP_MIN_LOG_LEVEL=0:

>>> import jax
>>> jax.devices()
2021-12-03 14:10:23.408083: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:171] XLA service 0x2d12d00 initialized for platform Interpreter (this does not guarantee that XLA will be used). Devices:
2021-12-03 14:10:23.408126: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:179] StreamExecutor device (0): Interpreter, <undefined>
2021-12-03 14:10:23.419007: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:171] XLA service 0x2cd8ea0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2021-12-03 14:10:23.419039: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:179] StreamExecutor device (0): Host, Default Version
2021-12-03 14:10:23.475845: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/tpu_client.cc:222] TPU platform initialization failed: Invalid argument: Invalid fd: -1; Couldn't open device: /dev/accel0 (Operation not permitted); Unable to create Node RegisterInterface for node 0, config: device_path: "/dev/accel0" mode: KERNEL debug_data_directory: "" dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true; could not create driver instance
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]

I expected it to still find all the TPUs, and want to run a script on one TPU at a time.

In a more recent version of jax: jax-0.2.25 and jaxlib-0.1.73, the error is more descriptive (which made me realise this is known, possibly not a bug per-se):

>>> import jax
WARNING: Logging before InitGoogle() is written to STDERR
I1203 14:15:18.362570 117716 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Not attempting to load libtpu.so in this process.
>>> jax.devices()
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[<jaxlib.xla_extension.Device object at 0x7f042db18570>]

It seems libtpu.so cannot be used in two separate python processes importing jax, and it would be helpful if this were possible?