I'm curious, have you considered using PyTorch or JAX for tensor processing? ML libraries seem to be much further along when it comes to performing compute-intensive, hardware-accelerated operations on Tensor's. And you get gradients basically for free (in terms of developer time). Also, the kernel compiler being added PyTorch 2 looks very promising.