2024年11月21日
Jax中关于NonZero的使用
摘要： Jax中关于NonZero的使用 在使用Jax的过程中，有时候会遇到函数输出是一个动态的Shape，这种情况下我们很难利用到Jax的即时编译的功能，不能使得性能最大化。这也是使用Tensor数据结构来计算的一个特点，有好有坏。本文介绍了Jax的另外一个函数NonZero，可以使得我们能够编译那些动态Shape输出的函数。 阅读全文
posted @ 2024-11-21 11:15 DECHIN 阅读(147) 评论(0) 推荐(0)