  void Compute(OpKernelContext* context) override {
    // The second args is N, which must be a positive scalar.
    const auto& n_in = context->input(1);
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(n_in.shape()),
        errors::InvalidArgument("N must be scalar but has rank ", n_in.dims()));
    int n = n_in.scalar<int32>()();
    OP_REQUIRES(context, n >= 0,
                errors::InvalidArgument("n must be non-negative but is ", n));

    // The first args is input tensor, which must have 1 dimension at least.
    const Tensor& input_in = context->input(0);
    const int num_dims = input_in.dims();
    OP_REQUIRES(context, num_dims >= 1,
                errors::InvalidArgument(
                    "Input must be at least rank 1 but is rank ", num_dims));
    // The last dimension of input tensor must be greater than N.
    OP_REQUIRES(
        context, input_in.dim_size(num_dims - 1) > n,
        errors::InvalidArgument("Input must have last dimension > n = ", n));

    // std::nth_element only support the nth-smallest selection.
    if (reverse_) {
      n = input_in.dim_size(num_dims - 1) - n - 1;
    }

    // Assume input_shape is [d1,d2,...dk], and output_shape is [d1,d2...dk-1].
    TensorShape out_shape;
    for (int i = 0; i < num_dims - 1; ++i) {
      out_shape.AddDim(input_in.dim_size(i));
    }
    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, out_shape, &output_tensor));

    functor::NthElementFunctor<Device, T> nthElementFunc;
    nthElementFunc(context, input_in, *output_tensor, n, reverse_);
  }
  void Compute(OpKernelContext* context) override {
     
    const auto& n_in = context->input(1);
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(n_in.shape()),
        errors::InvalidArgument("N must be scalar but has rank ", n_in.dims()));
    int n = n_in.scalar<int32>()();
    OP_REQUIRES(context, n >= 0,
                errors::InvalidArgument("n must be non-negative but is ", n));

     
    const Tensor& input_in = context->input(0);
    const int num_dims = input_in.dims();
    OP_REQUIRES(context, num_dims >= 1,
                errors::InvalidArgument(
                    "Input must be at least rank 1 but is rank ", num_dims));
     
    OP_REQUIRES(
        context, input_in.dim_size(num_dims - 1) > n,
        errors::InvalidArgument("Input must have last dimension > n = ", n));

     
    if (reverse_) {
      n = input_in.dim_size(num_dims - 1) - n - 1;
    }

     
    TensorShape out_shape;
    for (int i = 0; i < num_dims - 1; ++i) {
      out_shape.AddDim(input_in.dim_size(i));
    }
    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, out_shape, &output_tensor));

    functor::NthElementFunctor<Device, T> nthElementFunc;
    nthElementFunc(context, input_in, *output_tensor, n, reverse_);
  }