[GitHub] [tvm] zhuzilin edited a comment on pull request #8056: [Relay, TOPI] Add negative log likelihood loss (nll_loss) op

[GitBox <gi...@apache.org>]

zhuzilin edited a comment on pull request #8056:
URL: https://github.com/apache/tvm/pull/8056#issuecomment-868179165


   > in this case we need to change tag to `kInjective` as the reduction op is not broadcast
   
   @vinx13 Changing tag to `kInjective` will fail the tag check in `traverse_after_reduce` in
   https://github.com/apache/tvm/blob/d0791d3db971a111826d96201bd1e4c9c0d531da/python/tvm/topi/x86/reduction.py#L94
   whereas the empty tag was triggering `traverse_before_reduce`....
   
   It's worth noticing that the tag is of:
   
   ```c++
     auto T = tvm::te::compute(
         targets->shape,
         ...
         name, tag);
   ```
   
   which is an element-wise operation on `targets`. And when `reduction="mean"` or `reduction="sum"`, `T` will be reduced, whereas when `reduction="none"`, the `T` will be returned as the result of the `nll_loss`.
   
   Therefore, as `nll_loss` is registered to use reduce schedule, when `reduction="mean"` or `"sum"`, the `T` will go through the checks in `traverse_before_reduce` while when `reduction="none"`, `T` will go through checks in `traverse_after_reduce`. Right now, the `kBroadcast` happens to be the only tag to satisfiy both checks and passed the CI . To really solve this problem, I'm afraid we may need to define different tag and different scheduling for different op attr...


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org