I have wanted a JAX-like (XLA backed) Rust tensor library for ever!

Plus it looks like they have support for dynamic shapes: https://tensor4all.org/tenferro-rs/design/dynamic-symbolic-s...