| 30 | explicit SnapshotOp(OpKernelConstruction* context) : OpKernel(context) {} |
| 31 | |
| 32 | void Compute(OpKernelContext* context) override { |
| 33 | const Tensor& input = context->input(0); |
| 34 | Tensor* output = nullptr; |
| 35 | // Try to use buffer forwarding to avoid an explicit copy. |
| 36 | OP_REQUIRES_OK(context, context->forward_input_or_allocate_output( |
| 37 | {0}, 0, input.shape(), &output)); |
| 38 | if (!output->SharesBufferWith(input)) { |
| 39 | functor::Snapshot<Device, Scalar> functor; |
| 40 | functor(context->eigen_device<Device>(), input.flat<Scalar>(), |
| 41 | output->flat<Scalar>()); |
| 42 | } |
| 43 | } |
| 44 | }; |
| 45 | |
| 46 | #define REGISTER_KERNEL(TYPE) \ |
nothing calls this directly
no test coverage detected