[PATCH] D153422: [mlir][Linalg] Add a softmax op

[Renato Golin via Phabricator via llvm-commits]

rengolin added inline comments.


================
Comment at: mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td:100
+    For a given input tensor and a specified dimension `d`, compute:
+      1. the max `m` along that dimension `d`
+      2. f(x) = exp(x - m)
----------------
Would it make sense to allow more than one dimension here?

If your tensor is `MB x Head x TileX x TileY` you want the `max` over `TileX x TileY`, no?


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D153422/new/

https://reviews.llvm.org/D153422