| 92 | explicit BincountOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} |
| 93 | |
| 94 | void Compute(OpKernelContext* ctx) override { |
| 95 | const Tensor& arr_t = ctx->input(0); |
| 96 | const Tensor& size_tensor = ctx->input(1); |
| 97 | OP_REQUIRES(ctx, size_tensor.dims() == 0, |
| 98 | errors::InvalidArgument("Shape must be rank 0 but is rank ", |
| 99 | size_tensor.dims())); |
| 100 | const Tensor& weights_t = ctx->input(2); |
| 101 | |
| 102 | int32 size = size_tensor.scalar<int32>()(); |
| 103 | OP_REQUIRES( |
| 104 | ctx, size >= 0, |
| 105 | errors::InvalidArgument("size (", size, ") must be non-negative")); |
| 106 | |
| 107 | const auto arr = arr_t.flat<int32>(); |
| 108 | const auto weights = weights_t.flat<T>(); |
| 109 | Tensor* output_t; |
| 110 | OP_REQUIRES_OK(ctx, |
| 111 | ctx->allocate_output(0, TensorShape({size}), &output_t)); |
| 112 | auto output = output_t->flat<T>(); |
| 113 | OP_REQUIRES_OK(ctx, functor::BincountFunctor<Device, T>::Compute( |
| 114 | ctx, arr, weights, output)); |
| 115 | } |
| 116 | }; |
| 117 | |
| 118 | #define REGISTER_KERNELS(type) \ |
nothing calls this directly
no test coverage detected