| 180 | } |
| 181 | |
| 182 | xla::StatusOr<std::vector<InputCoords>> GetComputationInputs( |
| 183 | OpKernelContext* context, const char* input_name) { |
| 184 | OpInputList arg_list; |
| 185 | TF_RETURN_IF_ERROR(context->input_list(input_name, &arg_list)); |
| 186 | // Concatenate all input uids from list of scalars-or-vectors carrying them. |
| 187 | std::vector<InputCoords> input_coords; |
| 188 | for (int i = 0; i < arg_list.size(); ++i) { |
| 189 | const Tensor& arg = arg_list[i]; |
| 190 | if (TensorShapeUtils::IsScalar(arg.shape())) { |
| 191 | input_coords.emplace_back(arg.scalar<int64>()()); |
| 192 | } else { |
| 193 | TF_RET_CHECK(TensorShapeUtils::IsVector(arg.shape())); |
| 194 | auto arg_vec = arg.vec<int64>(); |
| 195 | const int64 num_elts = arg.shape().dim_size(0); |
| 196 | for (int i = 0; i < num_elts; ++i) { |
| 197 | input_coords.emplace_back(arg_vec(i)); |
| 198 | } |
| 199 | } |
| 200 | } |
| 201 | return std::move(input_coords); |
| 202 | } |
| 203 | |
| 204 | Status CreateExecuteOutput(OpKernelContext* context, |
| 205 | XRTMemoryManager* memory_manager, |
no test coverage detected