Builds XLA computations for each of the arguments to the computation. `args` are the arguments to the computation.
| 800 | // Builds XLA computations for each of the arguments to the computation. |
| 801 | // `args` are the arguments to the computation. |
| 802 | Status XlaCompiler::BuildArguments( |
| 803 | const Graph& graph, const std::vector<XlaCompiler::Argument>& args, |
| 804 | bool use_tuple_arg, xla::XlaBuilder* builder, XlaContext* context, |
| 805 | const std::map<int, xla::OpSharding>& arg_shardings, |
| 806 | std::vector<XlaExpression>* arg_expressions, |
| 807 | std::vector<int>* input_to_args, std::vector<xla::Shape>* input_shapes, |
| 808 | bool is_entry_computation) { |
| 809 | arg_expressions->resize(args.size()); |
| 810 | |
| 811 | // Argument numbers of arguments and resources that are to be passed to the |
| 812 | // XLA computation as runtime parameters. `input_to_args[a] = b` means that |
| 813 | // the a'th XLA input corresponds to the b'th original arg indexes. |
| 814 | input_to_args->clear(); |
| 815 | input_to_args->reserve(args.size()); |
| 816 | |
| 817 | // Fills in constant arguments, and computes non-constant argument order. |
| 818 | for (std::vector<XlaCompiler::Argument>::size_type i = 0; i < args.size(); |
| 819 | ++i) { |
| 820 | const XlaCompiler::Argument& arg = args[i]; |
| 821 | XlaExpression& arg_expression = (*arg_expressions)[i]; |
| 822 | switch (arg.kind) { |
| 823 | case XlaCompiler::Argument::kResource: { |
| 824 | TF_RET_CHECK(arg.resource_kind != XlaResource::kInvalid); |
| 825 | TF_RET_CHECK(absl::holds_alternative<TensorShape>(arg.shape)); |
| 826 | // TODO(phawkins): this code assumes that resource arguments do not |
| 827 | // alias. |
| 828 | XlaResource* resource = |
| 829 | context->AddResource(absl::make_unique<XlaResource>( |
| 830 | arg.resource_kind, i, arg.name, arg.type, |
| 831 | absl::get<TensorShape>(arg.shape), xla::XlaOp(), |
| 832 | /*max_array_size=*/arg.max_array_size, |
| 833 | /*tensor_array_gradients=*/arg.tensor_array_gradients, |
| 834 | /*tensor_array_multiple_writes_aggregate=*/true)); |
| 835 | arg_expression = XlaExpression::Resource(resource); |
| 836 | if (arg.initialized) { |
| 837 | input_to_args->push_back(i); |
| 838 | } |
| 839 | break; |
| 840 | } |
| 841 | case XlaCompiler::Argument::kParameter: |
| 842 | case XlaCompiler::Argument::kTensorList: |
| 843 | case XlaCompiler::Argument::kToken: { |
| 844 | input_to_args->push_back(i); |
| 845 | break; |
| 846 | } |
| 847 | case XlaCompiler::Argument::kConstant: |
| 848 | arg_expression = XlaExpression::Constant(arg.constant_value); |
| 849 | break; |
| 850 | case XlaCompiler::Argument::kInvalid: |
| 851 | return errors::Internal( |
| 852 | "Unreachable case in BuildArguments() while filling constant args"); |
| 853 | } |
| 854 | } |
| 855 | |
| 856 | if (input_to_args->empty() && !use_tuple_arg) { |
| 857 | return Status::OK(); |
| 858 | } |
| 859 |
nothing calls this directly
no test coverage detected