// CHECK: return %[[RESULT]]
  %0 = "tf.LogSoftmax"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32>
  func.return %0: tensor<2x3xf32>
}

// CHECK-LABEL: func @unranked_logsoftmax
func.func @unranked_logsoftmax(%arg0: tensor<*xf32>) -> tensor<*xf32> {
  // CHECK-NOT: "tf.LogSoftmax"
  %0 = "tf.LogSoftmax"(%arg0) : (tensor<*xf32>) -> tensor<*xf32>
  func.return %0: tensor<*xf32>
}

          {"PWRED",
           {"ArgMax", "ArgMin", "DiagPart", "Softmax",
            "SparseSoftmaxCrossEntropyWithLogits", "LogSoftmax"}},
          {"REDUCEWINDOW",
           {"ArgMax", "ArgMin", "DiagPart", "Softmax",
            "SparseSoftmaxCrossEntropyWithLogits", "LogSoftmax"}},
          {"REDUCEWINDOWPW", {"BiasAddGrad", "LRN", "LRNGrad"}},
          {"BN",

          concat_dim);
      result = concat_op.getResult();
    }

    rewriter.replaceOp(op, result);
    return success();
  }
};

// Decomposes Softmax and LogSoftmax to primitive TF ops, using the following
// formulas:
//
//     softmax = div(exp(logits), sum(exp(logits)))
//     log_softmax = sub(logits, log(sum(exp(logits))))
//