The Nx architecture is quite good. I started work on a backend that avoids requiring jvm/bazel/python toolchain (substituting in a zig toolchain):

https://github.com/ityonemo/ez

the two languages feel like they were made for each other. Despite the verbosity of doing some things in zig (the dynamic tensor broadcast function is a bit of a beast), due to both language's similar compile-time metaprogramming facilities, getting to the point of plugging in Nx's tensor addition + tensor multiplication is about 200 LoC.