| 33 | explicit RangeOp(OpKernelConstruction* context) : OpKernel(context) {} |
| 34 | |
| 35 | void Compute(OpKernelContext* context) override { |
| 36 | const Tensor& start_in = context->input(0); |
| 37 | const Tensor& limit_in = context->input(1); |
| 38 | const Tensor& delta_in = context->input(2); |
| 39 | OP_REQUIRES(context, IsLegacyScalar(start_in.shape()), |
| 40 | errors::InvalidArgument("start must be a scalar, not shape ", |
| 41 | start_in.shape().DebugString())); |
| 42 | OP_REQUIRES(context, IsLegacyScalar(limit_in.shape()), |
| 43 | errors::InvalidArgument("limit must be a scalar, not shape ", |
| 44 | limit_in.shape().DebugString())); |
| 45 | OP_REQUIRES(context, IsLegacyScalar(delta_in.shape()), |
| 46 | errors::InvalidArgument("delta must be a scalar, not shape ", |
| 47 | delta_in.shape().DebugString())); |
| 48 | const T start = start_in.scalar<T>()(); |
| 49 | const T limit = limit_in.scalar<T>()(); |
| 50 | const T delta = delta_in.scalar<T>()(); |
| 51 | OP_REQUIRES(context, delta != 0, |
| 52 | errors::InvalidArgument("Requires delta != 0: ", delta)); |
| 53 | if (delta > 0) { |
| 54 | OP_REQUIRES( |
| 55 | context, start <= limit, |
| 56 | errors::InvalidArgument( |
| 57 | "Requires start <= limit when delta > 0: ", start, "/", limit)); |
| 58 | } else { |
| 59 | OP_REQUIRES( |
| 60 | context, start >= limit, |
| 61 | errors::InvalidArgument( |
| 62 | "Requires start >= limit when delta < 0: ", start, "/", limit)); |
| 63 | } |
| 64 | auto size_auto = (std::is_integral<T>::value |
| 65 | ? (Eigen::numext::abs(limit - start) + |
| 66 | Eigen::numext::abs(delta) - T(1)) / |
| 67 | Eigen::numext::abs(delta) |
| 68 | : Eigen::numext::ceil( |
| 69 | Eigen::numext::abs((limit - start) / delta))); |
| 70 | OP_REQUIRES( |
| 71 | context, size_auto <= std::numeric_limits<int64>::max(), |
| 72 | errors::InvalidArgument("Requires ((limit - start) / delta) <= ", |
| 73 | std::numeric_limits<int64>::max())); |
| 74 | |
| 75 | int64 size = static_cast<int64>(size_auto); |
| 76 | |
| 77 | TensorShape shape; |
| 78 | OP_REQUIRES_OK(context, shape.AddDimWithStatus(size)); |
| 79 | Tensor* out = nullptr; |
| 80 | OP_REQUIRES_OK(context, |
| 81 | context->allocate_output(0, shape, &out)); |
| 82 | auto flat = out->flat<T>(); |
| 83 | T val = start; |
| 84 | for (int64 i = 0; i < size; ++i) { |
| 85 | flat(i) = T(val); |
| 86 | val += delta; |
| 87 | } |
| 88 | } |
| 89 | }; |
| 90 | |
| 91 | #define REGISTER_KERNEL(DEV, TYPE) \ |
nothing calls this directly
no test coverage detected