| 53 | |
| 54 | namespace { |
| 55 | Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph, |
| 56 | const std::vector<const XlaExpression*>& expressions, |
| 57 | std::vector<XlaCompiler::Argument>* args) { |
| 58 | auto client = ctx->compiler()->client(); |
| 59 | std::vector<bool> arg_must_be_compile_time_constant(expressions.size()); |
| 60 | |
| 61 | TF_RETURN_IF_ERROR(BackwardsConstAnalysis( |
| 62 | *graph, &arg_must_be_compile_time_constant, |
| 63 | /*compile_time_const_nodes=*/nullptr, ctx->function_library())); |
| 64 | |
| 65 | args->resize(expressions.size()); |
| 66 | for (int i = 0; i < args->size(); ++i) { |
| 67 | XlaCompiler::Argument& arg = (*args)[i]; |
| 68 | arg.type = ctx->input_type(i); |
| 69 | arg.shape = ctx->InputShape(i); |
| 70 | |
| 71 | switch (expressions[i]->kind()) { |
| 72 | case XlaExpression::Kind::kConstant: |
| 73 | arg.kind = XlaCompiler::Argument::kConstant; |
| 74 | arg.constant_value = expressions[i]->constant_value(); |
| 75 | break; |
| 76 | case XlaExpression::Kind::kXlaOp: |
| 77 | if (arg_must_be_compile_time_constant[i]) { |
| 78 | TF_ASSIGN_OR_RETURN(absl::optional<Tensor> value, |
| 79 | expressions[i]->ResolveConstant(client)); |
| 80 | if (!value.has_value()) { |
| 81 | return errors::InvalidArgument( |
| 82 | "Argument to function must be a compile-time constant, but " |
| 83 | "unable to resolve argument value to a constant."); |
| 84 | } |
| 85 | arg.kind = XlaCompiler::Argument::kConstant; |
| 86 | arg.constant_value = *value; |
| 87 | } else { |
| 88 | arg.kind = XlaCompiler::Argument::kParameter; |
| 89 | } |
| 90 | break; |
| 91 | case XlaExpression::Kind::kResource: { |
| 92 | XlaResource* resource = expressions[i]->resource(); |
| 93 | |
| 94 | arg.initialized = resource->initialized(); |
| 95 | arg.kind = XlaCompiler::Argument::kResource; |
| 96 | arg.resource_kind = resource->kind(); |
| 97 | arg.type = resource->type(); |
| 98 | arg.shape = resource->shape(); |
| 99 | arg.max_array_size = resource->max_array_size(); |
| 100 | arg.name = resource->name(); |
| 101 | break; |
| 102 | } |
| 103 | case XlaExpression::Kind::kTensorList: { |
| 104 | arg.kind = XlaCompiler::Argument::kTensorList; |
| 105 | const xla::XlaOp& tensor_list = expressions[i]->handle(); |
| 106 | arg.shape = tensor_list.builder()->GetShape(tensor_list).ValueOrDie(); |
| 107 | break; |
| 108 | } |
| 109 | case XlaExpression::Kind::kInvalid: |
| 110 | return errors::InvalidArgument("Invalid function argument"); |
| 111 | } |
| 112 | } |
no test coverage detected