| 27 | explicit AddNOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} |
| 28 | |
| 29 | void Compile(XlaOpKernelContext* ctx) override { |
| 30 | if (!ctx->ValidateInputsAreSameShape(this)) return; |
| 31 | |
| 32 | OP_REQUIRES(ctx, ctx->num_inputs() >= 1, |
| 33 | errors::InvalidArgument("AddN requires at least one argument")); |
| 34 | |
| 35 | XlaExpression::Kind kind = ctx->InputExpression(0).kind(); |
| 36 | xla::XlaOp sum; |
| 37 | switch (kind) { |
| 38 | case XlaExpression::Kind::kTensorList: { |
| 39 | // Check that all TensorLists are initialized. |
| 40 | for (int i = 1; i < ctx->num_inputs(); ++i) { |
| 41 | xla::XlaOp list = ctx->Input(i); |
| 42 | bool is_initialized; |
| 43 | OP_REQUIRES_OK(ctx, IsTensorListInitialized(list, &is_initialized)); |
| 44 | OP_REQUIRES( |
| 45 | ctx, is_initialized, |
| 46 | errors::InvalidArgument("TensorList input #", i, |
| 47 | " for AddN op is an uninitialized list")); |
| 48 | } |
| 49 | // Nested TensorList is not supported. |
| 50 | bool is_nested_list; |
| 51 | OP_REQUIRES_OK(ctx, IsNestedTensorList(ctx->Input(0), &is_nested_list)); |
| 52 | OP_REQUIRES(ctx, !is_nested_list, |
| 53 | errors::Unimplemented( |
| 54 | "Nested TensorList is not supported for AddN op")); |
| 55 | |
| 56 | OP_REQUIRES_OK(ctx, GetTensorListBuffer(ctx->Input(0), &sum)); |
| 57 | xla::Shape sum_shape; |
| 58 | OP_REQUIRES_OK(ctx, |
| 59 | GetTensorListBufferShape(ctx->Input(0), &sum_shape)); |
| 60 | for (int i = 1; i < ctx->num_inputs(); ++i) { |
| 61 | xla::XlaOp operand; |
| 62 | OP_REQUIRES_OK(ctx, GetTensorListBuffer(ctx->Input(i), &operand)); |
| 63 | // Check that the shapes match. |
| 64 | xla::Shape operand_shape; |
| 65 | OP_REQUIRES_OK( |
| 66 | ctx, GetTensorListBufferShape(ctx->Input(i), &operand_shape)); |
| 67 | OP_REQUIRES( |
| 68 | ctx, sum_shape.dimensions() == operand_shape.dimensions(), |
| 69 | errors::InvalidArgument( |
| 70 | "TensorList arguments to AddN must all have the same ", |
| 71 | "shape.\n", "Expected: ", sum_shape.DebugString(), "\n", |
| 72 | "Found: ", operand_shape.DebugString())); |
| 73 | sum = xla::Add(sum, operand); |
| 74 | } |
| 75 | xla::XlaOp push_index; |
| 76 | OP_REQUIRES_OK(ctx, GetTensorListPushIndex(ctx->Input(0), &push_index)); |
| 77 | OP_REQUIRES_OK(ctx, BuildNonNestedTensorList(sum, push_index, &sum)); |
| 78 | ctx->SetTensorListOutput(0, sum); |
| 79 | break; |
| 80 | } |
| 81 | default: |
| 82 | sum = ctx->Input(0); |
| 83 | for (int i = 1; i < ctx->num_inputs(); ++i) { |
| 84 | sum = xla::Add(sum, ctx->Input(i)); |
| 85 | } |
| 86 | ctx->SetOutput(0, sum); |
nothing calls this directly
no test coverage detected