DEV: pin CUDA variant for PyTorch

[crusaderky]

No description provided.

[lucascolley]

@traversaro could I double check that this is correct?

[traversaro]

This is a bit tricky. It is correct that having system-requirements = { cuda = "12" } and jax and pytorch as dependencies will result in CUDA-enabled jaxlib and pytorch being installed, as in both cases the CUDA variant have higher priority (see https://github.com/conda-forge/pytorch-cpu-feedstock/blob/46274b8459ee640a9f90a90d75ac931b770673ed/recipe/meta.yaml#L5-L10 and https://github.com/conda-forge/jaxlib-feedstock/blob/8d785af5387be376037f0210fc6dc2f5da95613c/recipe/meta.yaml#L4-L6). However, have a constraint on the build string ensures that this will always be true in the future.

Let's make an example. If you remove the jaxlib = { version = "*", build = "cuda*" } constraint, and let's pretend that a new jaxlib 0.100.0 (this is a made up number) is released. Let's also pretend that jaxlib 0.100.0 was released without cuda-version==12 compatible builds, for example as cuda 12 was dropped and only cuda 13 is supported, or as there was a regression in cuda builds for 0.100.0, that were marked as broken while cpu builds were still available. In that case, if you do not have jaxlib = { version = "*", build = "cuda*" } , if you refresh the lockfile the solver will just silently install a cpu-only version of jaxlib, while an error will be printed if jaxlib = { version = "*", build = "cuda*" } is present. In a nutshell I find having jaxlib = { version = "*", build = "cuda*" } a bit more robust, even if someone could disagree as it is not super-clear if build string structure is part of the public interface of the conda packages. The long term clean solution for this problem is probably flags (see conda/ceps#111 and https://prefix.dev/blog/ceps_2025), but that still needs to be discussed at the CEP-level.

[crusaderky]

for example as cuda 12 was dropped and only cuda 13 is supported [...] the solver will just silently install a cpu-only version of jaxlib

This is a strong argument against this PR. I'll make it explicit again.

[crusaderky]

@lucascolley ready for review again

[lucascolley]

thanks for the fast help as always @traversaro !! much appreciated

[lucascolley]

thanks Guido