[Feature Request] Running XLA functions on TPU

[frasermince]

Describe the bug

Upon running recv from the xla functions with jax tpus enabled I get the error: "NotImplementedError: MLIR translation rule for primitive 'AtariGymEnvPool_140079244014352_recv' not found for platform tpu". This is happening in a google tpu vm with python version 3.9.13, jax version 0.3.13, and envpool version 0.6.2.

To Reproduce

import jax
import envpool

env = envpool.make("Pong-v5", env_type="gym", num_envs=64, batch_size=16)
handle, recv, send, _ = env.xla()

handle2, data = recv(handle)

Expected behavior

I expect recv to work correctly on a tpu in a jitted jax program

Screenshots

System info

Describe the characteristic of your environment:

• Python 3.9.13
• Jax 0.3.13
• Envpool 0.6.2
• On a google TPU VM v3-8

Checklist

• [x] I have checked that there is no similar issue in the repo (required)
• [x] I have read the documentation (required)
• [x] I have provided a minimal working example to reproduce the bug (required)

[Trinkle23897]

@mavenlin could you please take a look?

[mavenlin]

There's currently no open API for copying between cpu and tpu as far as I know. Therefore not possible to support jitting on tpu for now.

[frasermince]

I see. So would it be possible to use the xla functions on the CPU as part of a jitted function and then move arrays over to the tpu afterwards? I've tried setting the default device to cpu while calling recv but I get the same error.

[mavenlin]

I see. So would it be possible to use the xla functions on the CPU as part of a jitted function and then move arrays over to the tpu afterwards? I've tried setting the default device to cpu while calling recv but I get the same error.

Do you mean that even jitting for CPU fails on a tpu vm? Before we introduced GPU support, jitting for CPU works fine on a GPU machine.

[qlan3]

I encountered the same problem.