Comment by phronimos

Numba is a great option for speeding up (vectorizing) loops and NumPy code, apart from CuPy and JAX. Xarray is also worth trying for tensors beyond 2 dimensions.