使用TensorFlow做矩阵计算是否可行？相比CuPy+Numba有何优劣？

• 完全可行。TensorFlow的底层核心本身就是通用张量计算框架，机器学习只是其上层主流应用场景之一，所有tf.linalg下的矩阵运算算子原生支持XLA编译、GPU/TPU加速能力，完全不需要绑定模型训练相关的逻辑。你不需要定义任何可训练变量、损失函数、优化器这类ML专属组件，只调用基础矩阵运算接口就能拿到对应的硬件加速收益，业内也有不少高性能计算场景会用这类ML框架做通用张量加速，不属于奇技淫巧。

• 最直接的方案是用tf.function装饰你的计算逻辑，同时开启jit_compile=True参数，TensorFlow会在你第一次运行该函数时完成计算图编译和XLA优化，后续所有调用都会直接复用编译好的计算图，不会重复编译。
• 如果需要跨进程/跨启动复用编译结果，可以开启TensorFlow的计算图缓存能力：设置环境变量TF_XLA_FLAGS=--xla_dump_to=/path/to/cache/dir --xla_load_compiled_modules_from_dump，就能把编译好的XLA二进制缓存到本地，下次启动直接加载，跳过编译步骤。
• 你的场景不需要保存训练参数，所以不用存完整的SavedModel，只需要把tf.function导出成轻量的ConcreteFunction保存即可，加载后直接调用就能用，完全不需要加载额外的训练相关数据。

• 学习成本低，API完全对齐NumPy，不需要了解TensorFlow的计算图、张量类型转换这类专属概念，现有NumPy写的矩阵计算代码几乎不用改就能换成CuPy跑在GPU上，Numba也只需要加装饰器就能给Python/NumPy代码加GPU加速，适配成本远低于TensorFlow。
• 无额外的冷启动编译开销：CuPy的算子都是提前编译好的，调用的时候直接执行，没有TensorFlow第一次运行的图编译等待时间，小批量、高频次调用的场景下延迟低很多。
• 灵活性更高：Numba支持自定义任意计算逻辑的GPU核函数，你可以针对自己的矩阵计算场景做定制化优化，不需要受限于TensorFlow已有的算子实现，也不需要写CUDA C++就能实现自定义高性能算子。
• 依赖更轻量：CuPy和Numba的安装包体积、运行时内存占用都远小于TensorFlow，部署的时候不需要带TensorFlow整套ML相关的依赖库，环境配置更简单。

内容的提问来源于stack exchange，提问作者material bug