func.func @softmax() {
  %input = util.unfoldable_constant dense<1.0> : tensor<2x8x4xf32>

  %init = tensor.empty() : tensor<2x8x4xf32>
  %1 = iree_linalg_ext.softmax dimension(2)
       ins(%input : tensor<2x8x4xf32>)
       outs(%init : tensor<2x8x4xf32>) -> tensor<2x8x4xf32>
  check.expect_almost_eq_const(
      %1,
      dense<0.25> : tensor<2x8x4xf32>
  ) : tensor<2x8x4xf32>
  return
}