> Imagine a world where machine learning models are written in Rust and can run on both Nvidia and AMD

Not likely in the next decade if ever. Unfortunately, the entire ecosystems of jax and torch are python based. Imagine retraining all those devs to use rust tooling.