| 41 | explicit TileOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} |
| 42 | |
| 43 | void Compile(XlaOpKernelContext* ctx) override { |
| 44 | const TensorShape input_shape = ctx->InputShape("input"); |
| 45 | const TensorShape multiples_shape = ctx->InputShape("multiples"); |
| 46 | |
| 47 | OP_REQUIRES( |
| 48 | ctx, TensorShapeUtils::IsVector(multiples_shape), |
| 49 | errors::InvalidArgument("Expected multiples to be 1-D, but got shape ", |
| 50 | multiples_shape.DebugString())); |
| 51 | OP_REQUIRES(ctx, input_shape.dims() == multiples_shape.num_elements(), |
| 52 | errors::InvalidArgument( |
| 53 | "Expected multiples argument to be a vector of length ", |
| 54 | input_shape.dims(), " but got length ", |
| 55 | multiples_shape.dim_size(0))); |
| 56 | const int input_dims = input_shape.dims(); |
| 57 | auto input = ctx->Input(0); |
| 58 | // If input is a scalar then multiples has 0 elements and this is |
| 59 | // a NoOp. |
| 60 | if (input_dims == 0) { |
| 61 | ctx->SetOutput(0, input); |
| 62 | return; |
| 63 | } |
| 64 | |
| 65 | std::vector<int64> multiples; |
| 66 | OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector("multiples", &multiples)); |
| 67 | std::vector<int64> output_dims(input_shape.dims()); |
| 68 | for (int64 i = 0; i < input_shape.dims(); ++i) { |
| 69 | OP_REQUIRES(ctx, multiples[i] >= 0, |
| 70 | errors::InvalidArgument("Expected multiples[", i, |
| 71 | "] >= 0, but got ", output_dims[i])); |
| 72 | output_dims[i] = input_shape.dim_size(i) * multiples[i]; |
| 73 | } |
| 74 | |
| 75 | // If all multiples are 1, than the input is the same as the output. |
| 76 | if (absl::c_all_of(multiples, |
| 77 | [](int64 multiple) { return multiple == 1; })) { |
| 78 | ctx->SetOutput(0, input); |
| 79 | return; |
| 80 | } |
| 81 | |
| 82 | bool can_tile_with_implicit_broadcast = true; |
| 83 | for (int i = 0; i < input_dims; ++i) { |
| 84 | int64 multiple = multiples[i]; |
| 85 | // If the multiple and input dimension are not 1, then tile cannot be |
| 86 | // implemented with a single hlo broadcast. |
| 87 | if (multiple != 1 && input_shape.dim_size(i) != 1) { |
| 88 | can_tile_with_implicit_broadcast = false; |
| 89 | } |
| 90 | } |
| 91 | |
| 92 | if (can_tile_with_implicit_broadcast) { |
| 93 | // Create a constant Zero the size of the output shape to leverage binary |
| 94 | // operation broadcast semantics. |
| 95 | auto broadcasted_zero = xla::Broadcast( |
| 96 | XlaHelpers::Zero(ctx->builder(), ctx->input_type(0)), output_dims); |
| 97 | if (ctx->input_type(0) == DT_BOOL) { |
| 98 | ctx->SetOutput(0, xla::Or(broadcasted_zero, input)); |
| 99 | } else { |
| 100 | ctx->SetOutput(0, xla::Add(broadcasted_zero, input)); |
nothing calls this directly
no test coverage detected