GPU - control the number of cores used

[AdrienCorenflos]

Hi,

Is there a way in the JAX config to control the capacity of the GPU used? I have this experiments where I want to reduce the parallelisation capacity of the GPU and see how the runtime of my function changes.

Adrien

[skye]

We don't expose such a config in JAX. Perhaps there's some generic way to do this? I'm gonna convert this issue to a discussion in case anyone knows of such a feature.