| func.func @softmax_static_10x256x256xf32() { |
| %input = util.unfoldable_constant dense<5.000000e+00> : tensor<10x256x256xf32> |
| %init = tensor.empty() : tensor<10x256x256xf32> |
| %0 = linalg.softmax dimension(2) ins(%input : tensor<10x256x256xf32>) outs(%init : tensor<10x256x256xf32>) -> tensor<10x256x256xf32> |
| check.expect_almost_eq_const(%0, dense<0.00390625> : tensor<10x256x256xf32>) : tensor<10x256x256xf32> |
| return |
| } |
| |
| func.func @softmax_dynamic_10x256x256xf32() { |
| %input = flow.tensor.dynamic_constant dense<5.000000e+00> : tensor<10x256x256xf32> -> tensor<?x?x?xf32> |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %d0 = tensor.dim %input, %c0 : tensor<?x?x?xf32> |
| %d1 = tensor.dim %input, %c1 : tensor<?x?x?xf32> |
| %d2 = tensor.dim %input, %c2 : tensor<?x?x?xf32> |
| %init = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32> |
| %0 = linalg.softmax dimension(2) ins(%input : tensor<?x?x?xf32>) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> |
| %barrier = util.optimization_barrier %0 : tensor<?x?x?xf32> |
| %result = tensor.cast %barrier : tensor<?x?x?xf32> to tensor<10x256x256xf32> |
| check.expect_almost_eq_const(%result, dense<0.00390625> : tensor<10x256x256xf32>) : tensor<10x256x256xf32> |
| return |
| } |
