Is it possible to create a hybrid mesh sharding of GPU and CPU?

[nasyxx]

I have tried

sharding = PositionalSharding(mesh_utils.create_device_mesh((3,), jax.devices() + jax.devices("cpu")))

However, JAX regarded the CPU as the GPU. At the same time, I couldn't put the array on it.

In hugging face, it is possible to put the model on GPU and CPU using device_map. Is it possible to do this in JAX?

[Kelvinlby]

+1

[yashk2810]

Not really. But if you describe what you are trying to do, I can probably suggest alternatives.