我在这里猜一猜。我也不能运行它，所以你可能要摆弄它

with jax.default_device(jax.devices("cpu")[0]):
    y = my_function()

看这里和这里的文件。

要直接指定执行函数的设备，请使用jax.jit的device参数。例如（使用GPU运行时，因为它是我目前可以访问的加速器）：

import jax

gpu_device = jax.devices('gpu')[0]
cpu_device = jax.devices('cpu')[0]

def my_function(x):
  return x.sum()

x = jax.numpy.arange(10)

x_gpu = jax.jit(my_function, device=gpu_device)(x)
print(x_tpu.device())
# gpu:0

x_cpu = jax.jit(my_function, device=cpu_device)(x)
print(x_cpu.device())
# TFRT_CPU_0

这也可以通过调用点周围的jax.default_device装饰器来控制：

with jax.default_device(cpu_device):
  print(jax.jit(my_function)(x).device())
  # TFRT_CPU_0

with jax.default_device(gpu_device):
  print(jax.jit(my_function)(x).device())
  # gpu:0