1 #pragma once 2 3 #include <torch/csrc/jit/tensorexpr/kernel.h> 4 5 namespace torch { 6 namespace jit { 7 namespace tensorexpr { 8 9 Tensor computeSoftmax( 10 const std::vector<ArgValue>& inputs, 11 const std::vector<ExprHandle>& outputShape, 12 const std::vector<ExprHandle>& outputStrides, 13 bool log_softmax); 14 15 } // namespace tensorexpr 16 } // namespace jit 17 } // namespace torch 18