using ops::Conv2DBackpropInput;
using ops::DepthwiseConv2dNative;
using ops::Elu;
using ops::FractionalAvgPool;
using ops::FractionalMaxPool;
using ops::FusedBatchNormV3;
using ops::L2Loss;
using ops::LogSoftmax;
using ops::LRN;
using ops::MaxPool;
using ops::MaxPool3D;
using ops::MaxPoolV2;
using ops::Placeholder;
using ops::Relu;
using ops::Relu6;
using ops::Selu;
using ops::Softmax;
using ops::Softplus;

    auto multiply_result = Multiply(scope, subtraction_result, logits_softmax);
    grad = Add(scope, grad, multiply_result);
  }
  auto minus_log_softmax = Multiply(scope, LogSoftmax(scope, logits), -1.0f);
  grad_outputs->push_back(grad);
  grad_outputs->push_back(BroadcastMul(scope, grad_loss, minus_log_softmax));
  return scope.status();
}
REGISTER_GRADIENT_OP("SoftmaxCrossEntropyWithLogits",

  // 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>
}

// computes loss and backprop of the loss with respect to 'features'.
//
// Softmax cross entropy loss is defined as follows:
//
//  loss = Sum(-labels * Log(Exp(features) / Sum(Exp(features)))
//  loss = Sum(-labels * LogSoftmax(features))
//
// Computing gradient of the loss with respect to features gives us,
//
//  backprop = (Exp(features) / Sum(Exp(features))) - labels
//  backprop = Softmax(features) - labels
//

          {"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))))
//

  %1 = "tf.Log"(%0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<8x16xf32>) -> tensor<8x16xf32>
  func.return %1: tensor<8x16xf32>

  // CHECK: %0 = "tf.LogSoftmax"(%arg0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<8x16xf32>) -> tensor<8x16xf32>
  // CHECK: return %0
}

// CHECK-LABEL: testLogToLog1p

  func.return %0 : tensor<8x16xf32>
// CHECK-LABEL: log
// CHECK:  "tfl.log"(%arg0) : (tensor<8x16xf32>) -> tensor<8x16xf32>
}

func.func @log_softmax(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
  %0 = "tf.LogSoftmax"(%arg0) : (tensor<8x16xf32>) -> tensor<8x16xf32>
  func.return %0 : tensor<8x16xf32>
// CHECK-LABEL: log_softmax
// CHECK:  "tfl.log_softmax"(%arg0) : (tensor<8x16xf32>) -> tensor<8x16xf32>
}

      return ArraysAreCastCompatible(inferred, actual);
    }
  }];
}

def TF_LogSoftmaxOp : TF_Op<"LogSoftmax", [Pure, TF_SameOperandsAndResultTypeResolveRef]> {
  let summary = "Computes log softmax activations.";

  let description = [{
For each batch `i` and class `j` we have

    logsoftmax[i, j] = logits[i, j] - log(sum(exp(logits[i])))
  }];

  let arguments = (ins

    stateless and do not touch any resources.
*   Refactors code in Quant8 LSTM support to reduce TFLite binary size.
*   Add support of local soft device placement for eager op.
*   Add HW acceleration support for `LogSoftMax`.
*   Added a function `nested_value_rowids` for ragged tensors.
*   Add guard to avoid acceleration of L2 Normalization with input rank != 4
*   Add `tf.math.cumulative_logsumexp operation`.
*   Add `tf.ragged.stack`.