如题，对于一维的张量data[lbk]N[rbk]，可以把这个维度分到cuda的block上去做，但是当张量是多维时这种做法就失效了。比如对于data[lbk]N1[rbk][lbk]N2[rbk][lbk]N3[rbk][lbk]...[rbk][lbk]Nk[rbk]，指定某个维度求max和argmax有没有高效的cuda实现？pytorch内部是怎么实现这个算子的呢？

可以找找源码里面的对应实现