| 255 | // Functor: defined in cwise_ops.h. E.g., functor::sqrt. |
| 256 | template <typename Device, typename Functor> |
| 257 | class UnaryOp : public OpKernel { |
| 258 | public: |
| 259 | typedef typename Functor::in_type Tin; // Input scalar data type. |
| 260 | typedef typename Functor::out_type Tout; // Output scalar data type. |
| 261 | // Tin may be different from Tout. E.g., abs: complex64 -> float |
| 262 | |
| 263 | explicit UnaryOp(OpKernelConstruction* ctx) : OpKernel(ctx) { |
| 264 | auto in = DataTypeToEnum<Tin>::v(); |
| 265 | auto out = DataTypeToEnum<Tout>::v(); |
| 266 | OP_REQUIRES_OK(ctx, ctx->MatchSignature({in}, {out})); |
| 267 | } |
| 268 | |
| 269 | void Compute(OpKernelContext* ctx) override { |
| 270 | const Tensor& inp = ctx->input(0); |
| 271 | Tensor* out = nullptr; |
| 272 | if (std::is_same<Tin, Tout>::value) { |
| 273 | OP_REQUIRES_OK(ctx, ctx->forward_input_or_allocate_output( |
| 274 | {0}, 0, inp.shape(), &out)); |
| 275 | } else { |
| 276 | OP_REQUIRES_OK(ctx, ctx->allocate_output(0, inp.shape(), &out)); |
| 277 | } |
| 278 | functor::UnaryFunctor<Device, Functor>()( |
| 279 | ctx->eigen_device<Device>(), out->flat<Tout>(), inp.flat<Tin>()); |
| 280 | } |
| 281 | }; |
| 282 | |
| 283 | template <typename Device, VariantUnaryOp OpEnum> |
| 284 | class UnaryVariantOp : public OpKernel { |