| 140 | explicit LinSpaceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} |
| 141 | |
| 142 | void Compile(XlaOpKernelContext* ctx) override { |
| 143 | const TensorShape start_in_shape = ctx->InputShape("start"); |
| 144 | const TensorShape stop_in_shape = ctx->InputShape("stop"); |
| 145 | const TensorShape num_in_shape = ctx->InputShape("num"); |
| 146 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(start_in_shape), |
| 147 | errors::InvalidArgument("start must be a scalar, not shape ", |
| 148 | start_in_shape.DebugString())); |
| 149 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(stop_in_shape), |
| 150 | errors::InvalidArgument("stop must be a scalar, not shape ", |
| 151 | stop_in_shape.DebugString())); |
| 152 | OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(num_in_shape), |
| 153 | errors::InvalidArgument("num must be a scalar, not shape ", |
| 154 | num_in_shape.DebugString())); |
| 155 | |
| 156 | DataType type = ctx->input_type(0); |
| 157 | |
| 158 | int64 num; |
| 159 | OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar("num", &num)); |
| 160 | OP_REQUIRES(ctx, num > 0, |
| 161 | errors::InvalidArgument("Requires num > 0: ", num)); |
| 162 | Tensor out_constant(type, TensorShape({num})); |
| 163 | |
| 164 | xla::Literal start_literal; |
| 165 | OP_REQUIRES_OK(ctx, ctx->ConstantInput("start", &start_literal)); |
| 166 | xla::Literal stop_literal; |
| 167 | OP_REQUIRES_OK(ctx, ctx->ConstantInput("stop", &stop_literal)); |
| 168 | |
| 169 | switch (type) { |
| 170 | case DT_FLOAT: { |
| 171 | float start = start_literal.GetFirstElement<float>(); |
| 172 | float stop = stop_literal.GetFirstElement<float>(); |
| 173 | auto flat = out_constant.flat<float>(); |
| 174 | if (num == 1) { |
| 175 | flat(0) = start; |
| 176 | } else { |
| 177 | const float step = (stop - start) / (num - 1); |
| 178 | for (int64 i = 0; i < num - 1; ++i) { |
| 179 | flat(i) = start + step * i; |
| 180 | } |
| 181 | // The last value in the sequence must be equal to stop. |
| 182 | flat(num - 1) = stop; |
| 183 | } |
| 184 | break; |
| 185 | } |
| 186 | case DT_DOUBLE: { |
| 187 | double start = start_literal.GetFirstElement<double>(); |
| 188 | double stop = stop_literal.GetFirstElement<double>(); |
| 189 | auto flat = out_constant.flat<double>(); |
| 190 | if (num == 1) { |
| 191 | flat(0) = start; |
| 192 | } else { |
| 193 | const double step = (stop - start) / (num - 1); |
| 194 | for (int64 i = 0; i < num - 1; ++i) { |
| 195 | flat(i) = start + step * i; |
| 196 | } |
| 197 | // The last value in the sequence must be equal to stop. |
| 198 | flat(num - 1) = stop; |
| 199 | } |
nothing calls this directly
no test coverage detected