| 62 | } |
| 63 | |
| 64 | void Compile(XlaOpKernelContext* ctx) override { |
| 65 | const TensorShape input_shape = ctx->InputShape("input"); |
| 66 | const TensorShape pad_shape = ctx->InputShape("paddings"); |
| 67 | |
| 68 | MirrorPadMode mode; |
| 69 | OP_REQUIRES_OK(ctx, GetNodeAttr(def(), "mode", &mode)); |
| 70 | OP_REQUIRES( |
| 71 | ctx, mode == MirrorPadMode::REFLECT || mode == MirrorPadMode::SYMMETRIC, |
| 72 | xla::Unimplemented("Unsupported MirrorPad mode. Only SYMMETRIC and " |
| 73 | "REFLECT modes are currently supported")); |
| 74 | |
| 75 | const int dims = input_shape.dims(); |
| 76 | OP_REQUIRES( |
| 77 | ctx, |
| 78 | TensorShapeUtils::IsMatrix(pad_shape) && pad_shape.dim_size(1) == 2, |
| 79 | errors::InvalidArgument("paddings must be a matrix with 2 columns: ", |
| 80 | pad_shape.DebugString())); |
| 81 | OP_REQUIRES( |
| 82 | ctx, dims == pad_shape.dim_size(0), |
| 83 | errors::InvalidArgument( |
| 84 | "The first dimension of paddings must be the rank of inputs", |
| 85 | pad_shape.DebugString(), " ", input_shape.DebugString())); |
| 86 | |
| 87 | // Evaluate the 'padding' constant input, reshaping to a matrix. |
| 88 | xla::Literal pad_literal; |
| 89 | OP_REQUIRES_OK(ctx, |
| 90 | ctx->ConstantInputAsInt64Literal("paddings", &pad_literal)); |
| 91 | |
| 92 | xla::XlaBuilder* b = ctx->builder(); |
| 93 | auto in0 = ctx->Input("input"); |
| 94 | xla::StatusOr<xla::Shape> in0_shape = b->GetShape(in0); |
| 95 | OP_REQUIRES(ctx, in0_shape.ok(), in0_shape.status()); |
| 96 | xla::StatusOr<xla::XlaOp> accum_status = |
| 97 | DoMirrorPad(in0, in0_shape.ValueOrDie(), pad_literal, mode, b); |
| 98 | |
| 99 | OP_REQUIRES_OK(ctx, accum_status.status()); |
| 100 | |
| 101 | ctx->SetOutput(0, accum_status.ValueOrDie()); |
| 102 | } |
| 103 | |
| 104 | private: |
| 105 | TF_DISALLOW_COPY_AND_ASSIGN(MirrorPadOp); |
nothing calls this directly
no test coverage detected