#include "tensorflow/c/experimental/ops/nn_ops.h"
#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
#include "tensorflow/core/platform/errors.h"

using std::vector;
using tensorflow::ops::BiasAddGrad;
using tensorflow::ops::ReluGrad;

namespace tensorflow {
namespace gradients {
namespace {

class ReluGradientFunction : public GradientFunction {
 public:
  explicit ReluGradientFunction(vector<AbstractTensorHandle*> f_outputs)

          (MHLO_ClampOp (MHLO_ConstantOp (GetScalarOfType<0> $input)), $input,
                       (MHLO_ConstantOp (GetScalarOfType<6> $input))),
          [(TF_IntOrFpTensor $input)]>;

// ReluGrad(gradients, features) = gradients * (features > 0)
// The condition that $gradients and $features need to have the same shape is
// implicitly enforced: $zero is created to have the same shape as $features,

            // Others
            "AddN", "Bitcast", "Cast", "ClipByValue", "Const", "Empty",
            "Identity", "IdentityN", "Relu", "Relu6", "ReluGrad", "Relu6Grad",
            "LeakyReluGrad", "Elu", "EluGrad", "Selu", "SeluGrad", "Select",
            "SelectV2", "Transpose", "ConjugateTranspose",
            "_UnaryOpsComposition", "CollectiveReduceV2",
            "CollectiveAssignGroupV2",

  func.func @testSameOperandsAndResultTypeResolveRefBinary(%lhs: tensor<2x3x?x?xf32>, %rhs: tensor<2x?x5x?xf32>) -> (tensor<?x?x?x?xf32>) {
    // CHECK: (tensor<2x3x?x?xf32>, tensor<2x?x5x?xf32>) -> tensor<2x3x5x?xf32>
    %0 = "tf.ReluGrad"(%lhs, %rhs) : (tensor<2x3x?x?xf32>, tensor<2x?x5x?xf32>) -> tensor<?x?x?x?xf32>
    func.return %0 : tensor<?x?x?x?xf32>
  }

  // CHECK-LABEL: func @test_xla_sharding

    static bool isCompatibleReturnTypes(TypeRange inferred, TypeRange actual) {
      return ArraysAreCastCompatible(inferred, actual);
    }
  }];
}

def TF_ReluGradOp : TF_Op<"ReluGrad", [Pure, TF_SameOperandsAndResultTypeResolveRef]> {
  let summary = "Computes rectified linear gradients for a Relu operation.";

  let arguments = (ins

    // CHECK-NEXT:  return %[[RES]] : tensor<4x8xf32>
    %2 = "tf.SeluGrad"(%gradients, %features) : (tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
    func.return %2 : tensor<4x8xf32>
}

// CHECK-LABEL: func @expm1
// CHECK-SAME: (%[[ARG0:.*]]: tensor<3x4xf32>)

  // CHECK: %[[RESULT:.*]] = mhlo.select %[[PRED]], %[[GRADIENTS]], %[[MULGRAD]]
  // CHECK: return %[[RESULT]]
  %2 = "tf.EluGrad"(%gradients, %features) : (tensor<4x8xf32>, tensor<?x?xf32>) -> tensor<4x8xf32>
  func.return %2 : tensor<4x8xf32>
}

//===----------------------------------------------------------------------===//