xla

0a1594a8 - Move mutable properties of env to thread local, misc changes (#9501)

Move mutable properties of env to thread local, misc changes (#9501) * Refactored jax device handling * Removed option to use CPU jax array for CPU torch tensors. - changing jax devices after the fact will use different APIs