Method Compute

tensorflow/core/kernels/count_up_to_op.cc:68–93

  }

  void Compute(OpKernelContext* context) override {
    core::RefCountPtr<Var> variable;
    OP_REQUIRES_OK(context, LookupResource(context, HandleFromInput(context, 0),
                                           &variable));
    mutex_lock l(*variable->mu());
    Tensor before_increment = *variable->tensor();
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(before_increment.shape()),
        errors::InvalidArgument("input is not a scalar: ",
                                before_increment.shape().DebugString()));
    if (before_increment.scalar<T>()() >= limit_) {
      context->SetStatus(errors::OutOfRange("Reached limit of ", limit_));
      return;
    }
    // Allocate new buffer
    AllocatorAttributes attr;
    attr.set_gpu_compatible(true);
    attr.set_nic_compatible(true);
    PersistentTensor unused;
    Tensor* tmp;
    OP_REQUIRES_OK(context, context->allocate_persistent(
                                dtype_, TensorShape({}), &unused, &tmp, attr));
    *variable->tensor() = *tmp;
    tmp->scalar<T>()() = before_increment.scalar<T>()() + 1;
    context->set_output(0, before_increment);
  }

 private:
  T limit_;