| 47 | explicit ArgOp(OpKernelConstruction* context) : OpKernel(context) {} |
| 48 | |
| 49 | void Compute(OpKernelContext* context) override { |
| 50 | const Tensor& input = context->input(0); |
| 51 | const Tensor& dimension = context->input(1); |
| 52 | |
| 53 | OP_REQUIRES(context, TensorShapeUtils::IsScalar(dimension.shape()), |
| 54 | errors::InvalidArgument( |
| 55 | "dim must be a scalar, but received tensor of shape: ", |
| 56 | dimension.shape().DebugString())); |
| 57 | |
| 58 | const int32 dim = internal::SubtleMustCopy(dimension.scalar<int32>()()); |
| 59 | const int input_dims = input.dims(); |
| 60 | |
| 61 | int axis = dim < 0 ? dim + input_dims : dim; |
| 62 | |
| 63 | OP_REQUIRES(context, FastBoundsCheck(axis, input_dims), |
| 64 | errors::InvalidArgument("Expected dimension in the range [", |
| 65 | -input_dims, ", ", input_dims, |
| 66 | "), but got ", dim)); |
| 67 | OP_REQUIRES( |
| 68 | context, input.dim_size(axis) > 0, |
| 69 | errors::InvalidArgument("Reduction axis ", dim, " is empty in shape ", |
| 70 | input.shape().DebugString())); |
| 71 | |
| 72 | TensorShape output_shape; |
| 73 | const TensorShape& input_shape = input.shape(); |
| 74 | for (int d = 0; d < input_dims - 1; ++d) { |
| 75 | output_shape.AddDim(input_shape.dim_size((d < axis) ? d : d + 1)); |
| 76 | } |
| 77 | Tensor* output = nullptr; |
| 78 | OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); |
| 79 | |
| 80 | if (output_shape.num_elements() == 0) { |
| 81 | return; |
| 82 | } |
| 83 | |
| 84 | #define HANDLE_DIM(NDIM) \ |
| 85 | case NDIM: \ |
| 86 | ArgFunctor::Reduce##NDIM(context->eigen_device<Device>(), \ |
| 87 | input.tensor<T, NDIM>(), axis, \ |
| 88 | output->tensor<Tout, NDIM - 1>()); \ |
| 89 | break; |
| 90 | |
| 91 | switch (input_dims) { |
| 92 | HANDLE_DIM(1); |
| 93 | HANDLE_DIM(2); |
| 94 | HANDLE_DIM(3); |
| 95 | HANDLE_DIM(4); |
| 96 | HANDLE_DIM(5); |
| 97 | |
| 98 | default: |
| 99 | OP_REQUIRES(context, false, |
| 100 | errors::InvalidArgument( |
| 101 | "ArgOp : Unhandled input dimensions: ", input_dims)); |
| 102 | } |
| 103 | } |
| 104 | #undef HANDLE_DIM |
| 105 | |
| 106 | private: |
nothing calls this directly
no test coverage detected