| 45 | */ |
| 46 | template <typename Device, typename T> |
| 47 | class ConditionalAccumulator |
| 48 | : public TypedConditionalAccumulatorBase<const Tensor> { |
| 49 | public: |
| 50 | // Args: |
| 51 | // dtype: The datatype of the gradients to be accumulated. |
| 52 | // shape: The shape of the accumulated gradients. |
| 53 | // name: A name to use for the ConditionalAccumulator. |
| 54 | // reduction_type: The reduction type, i.e., MEAN or SUM |
| 55 | ConditionalAccumulator(const DataType& dtype, const PartialTensorShape& shape, |
| 56 | const string& name, const string& reduction_type) |
| 57 | : TypedConditionalAccumulatorBase<const Tensor>(dtype, shape, name, |
| 58 | reduction_type) {} |
| 59 | ~ConditionalAccumulator() override{}; |
| 60 | |
| 61 | protected: |
| 62 | // accum_grad is the tensor that holds the aggregate gradient. |
| 63 | // It is initialized the first time ApplyGrad is called. |
| 64 | Tensor* accum_grad_ = nullptr; |
| 65 | PersistentTensor accum_grad_persistent_; |
| 66 | |
| 67 | functor::SetZeroFunctor<Device, T> set_zero_functor_; |
| 68 | |
| 69 | Status ValidateShape(const Tensor* tensor) |
| 70 | EXCLUSIVE_LOCKS_REQUIRED(this->mu_) { |
| 71 | // Must be compatible with accumulated gradient if available |
| 72 | if (counter_ > 0) { |
| 73 | if (!accum_grad_->shape().IsSameSize(tensor->shape())) { |
| 74 | return errors::InvalidArgument("Shape mismatch: expected ", |
| 75 | accum_grad_->shape().DebugString(), |
| 76 | ", got ", tensor->shape().DebugString()); |
| 77 | } |
| 78 | } |
| 79 | // Must also be compatible with given shape |
| 80 | if (!shape_.IsCompatibleWith(tensor->shape())) { |
| 81 | return errors::InvalidArgument("Shape mismatch: expected ", |
| 82 | shape_.DebugString(), ", got ", |
| 83 | tensor->shape().DebugString()); |
| 84 | } |
| 85 | return Status::OK(); |
| 86 | } |
| 87 | |
| 88 | void AllocateAndAssignToAccumGradFunction(OpKernelContext* ctx, |
| 89 | const Tensor* grad) override { |
| 90 | // TODO(b/32704451): Don't just ignore the ::tensorflow::Status object! |
| 91 | ctx->allocate_persistent(dtype_, grad->shape(), &accum_grad_persistent_, |
| 92 | &accum_grad_) |
| 93 | .IgnoreError(); |
| 94 | accum_grad_->flat<T>().device(ctx->template eigen_device<Device>()) = |
| 95 | grad->flat<T>(); |
| 96 | } |
| 97 | |
| 98 | void AddToAccumGradFunction(OpKernelContext* ctx, |
| 99 | const Tensor* grad) override { |
| 100 | accum_grad_->flat<T>().device(ctx->template eigen_device<Device>()) += |
| 101 | grad->flat<T>(); |
| 102 | } |
| 103 | |
| 104 | void DivideAccumGradByCounter(OpKernelContext* ctx) override |
nothing calls this directly
no test coverage detected