  static Status Compute(OpKernelContext* context,
                        const typename TTypes<T, 1>::ConstTensor& sorted_inputs,
                        const typename TTypes<T, 1>::ConstTensor& values,
                        int batch_size, int num_inputs, int num_values,
                        typename TTypes<OutType, 1>::Tensor* output) {
    // TODO(rmlarsen): add multithreading or interleaving.
    for (int b = 0; b < batch_size; ++b) {
      const T* sorted_inputs_ptr = sorted_inputs.data() + b * num_inputs;
      OutType* output_ptr = output->data() + b * num_values;
      for (int i = 0; i < num_values; ++i) {
        output_ptr[i] =
            std::upper_bound(sorted_inputs_ptr, sorted_inputs_ptr + num_inputs,
                             values(i + b * num_values)) -
            sorted_inputs_ptr;
      }
    }

    return Status::OK();
  }
  static Status Compute(OpKernelContext* context,
                        const typename TTypes<T, 1>::ConstTensor& sorted_inputs,
                        const typename TTypes<T, 1>::ConstTensor& values,
                        int batch_size, int num_inputs, int num_values,
                        typename TTypes<OutType, 1>::Tensor* output) {
     
    for (int b = 0; b < batch_size; ++b) {
      const T* sorted_inputs_ptr = sorted_inputs.data() + b * num_inputs;
      OutType* output_ptr = output->data() + b * num_values;
      for (int i = 0; i < num_values; ++i) {
        output_ptr[i] =
            std::upper_bound(sorted_inputs_ptr, sorted_inputs_ptr + num_inputs,
                             values(i + b * num_values)) -
            sorted_inputs_ptr;
      }
    }

    return Status::OK();
  }