Simplify Sparse Deep Learning with Universal Sparse Tensor in nvmath-python
In a previous post, we introduced the Universal Sparse Tensor (UST), enabling developers to decouple a tensor’s sparsity from its memory layout for greater flexibility and performance. We’re excited to announce the integration of the UST into nvmath-python v0.9.0 to accelerate sparse scientific and deep learning applications. This post provides a walkthrough of key UST features, implementation details, and performance overview, including: - Zero-cost interoperability: Data-movement-free conversion with PyTorch, SciPy, and CuPy. - Custom formats: Define novel sparsity schemes. - Polymorphic operations: Sparsity-agnostic functions automatically use optimized kernels or generate custom sparse code—eliminating the need for manual coding of new formats. - PyTorch injection: Easily inject UST performance benefits into existing PyTorch models. - Transparent caching: Avoid JIT/LTO recompilation and replanning—amortizing overhead over subsequent repeated execution of the same operation. Tensor format DSL The UST describes common (e.g., COO, CSR,…