| 35 | class CheckExecutionArityTest : public ClientLibraryTestBase {}; |
| 36 | |
| 37 | TEST_F(CheckExecutionArityTest, TwoParamComputationNumArguments) { |
| 38 | XlaBuilder builder("add_two_params"); |
| 39 | auto param_literal = LiteralUtil::CreateR1<float>({1.1f, 2.2f}); |
| 40 | |
| 41 | auto p0 = Parameter(&builder, 0, param_literal.shape(), "param0"); |
| 42 | auto p1 = Parameter(&builder, 1, param_literal.shape(), "param1"); |
| 43 | Add(p0, p1); |
| 44 | |
| 45 | auto param0_data = |
| 46 | client_->TransferToServer(param_literal).ConsumeValueOrDie(); |
| 47 | auto param1_data = |
| 48 | client_->TransferToServer(param_literal).ConsumeValueOrDie(); |
| 49 | |
| 50 | auto computation_status = builder.Build(); |
| 51 | ASSERT_IS_OK(computation_status.status()); |
| 52 | auto computation = computation_status.ConsumeValueOrDie(); |
| 53 | |
| 54 | // The arity of the UserComputation is 2 arguments. Execution will succeed |
| 55 | // with 2 arguments, but fail with a different number. |
| 56 | auto result_two_args = client_->Execute( |
| 57 | computation, {param0_data.get(), param1_data.get()}, &execution_options_); |
| 58 | ASSERT_IS_OK(result_two_args.status()); |
| 59 | |
| 60 | auto result_one_arg = |
| 61 | client_->Execute(computation, {param0_data.get()}, &execution_options_); |
| 62 | ASSERT_FALSE(result_one_arg.ok()); |
| 63 | ASSERT_EQ(result_one_arg.status().code(), |
| 64 | tensorflow::error::INVALID_ARGUMENT); |
| 65 | ASSERT_THAT(result_one_arg.status().error_message(), |
| 66 | ContainsRegex("takes 2")); |
| 67 | |
| 68 | auto result_zero_args = |
| 69 | client_->Execute(computation, {}, &execution_options_); |
| 70 | ASSERT_FALSE(result_zero_args.ok()); |
| 71 | ASSERT_EQ(result_zero_args.status().code(), |
| 72 | tensorflow::error::INVALID_ARGUMENT); |
| 73 | ASSERT_THAT(result_zero_args.status().error_message(), |
| 74 | ContainsRegex("takes 2")); |
| 75 | } |
| 76 | |
| 77 | XLA_TEST_F(CheckExecutionArityTest, CheckArgumentShapes) { |
| 78 | XlaBuilder builder("add_two_params"); |