How is env.xla() implemented?

[uduse]

From what I understand, only environments implemented in JAX can be JIT-compiled together with a JAX policy.
For example, gymnax implements a bunch of such environments in JAX.
However, it seems that in this project the non-JAX backend environments can be purified with env.xla().
I looked around in the code and it seems EnvPool is copying memory directly for some reason but I can't really figure out what we are doing here.
Can someone give a high-level explaination of how EnvPool implements its env.xla() interface? and how this interface is different from using environments implemented using native JAX?

[mavenlin]

Hi @uduse xla here is implemented via custom call, any c++ code (envpool code) can be put under a custom call and jitted into jax computation graph.

The GPU device support is faked via an extra call to copy the memory to GPU device.

https://dfm.io/posts/extending-jax/ this post is our main reference for the implementation.

[uduse]

This is very helpful! Thank you very much for the info 👍