Nvmath-Python: Nvidia Math Libraries for the Python Ecosystem
8 months ago
- #Python
- #NVIDIA
- #CUDA
- nvmath-python将NVIDIA数学库集成至Python环境，提供直观的API接口
- 支持多种执行空间，兼容Python数组/张量框架
- 可访问NVIDIA cuBLASLt库全部参数（部分参数在其他平台不可用）
- 使用示例包含矩阵乘法运算，支持混合精度计算类型等微调选项
- 支持设备端API调用，可在自定义设备内核中运行NVIDIA库函数
- 为FFT函数提供自定义前导/收尾程序，可将其编译为LTO-IR中间表示
- 示例展示通过自定义收尾程序实现FFT输出的单位缩放
- 基于Apache 2.0许可证，当前处于Beta测试阶段，欢迎提供改进建议