| 71 | explicit ReverseV2Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} |
| 72 | |
| 73 | void Compile(XlaOpKernelContext* ctx) override { |
| 74 | // r = tf.reverse(x, axes) |
| 75 | const TensorShape x_shape = ctx->InputShape(0); |
| 76 | const TensorShape axes_shape = ctx->InputShape(1); |
| 77 | // Validate input sizes. |
| 78 | OP_REQUIRES(ctx, TensorShapeUtils::IsVector(axes_shape), |
| 79 | errors::InvalidArgument("axes must be a vector, not shape ", |
| 80 | axes_shape.DebugString())); |
| 81 | OP_REQUIRES(ctx, axes_shape.num_elements() <= x_shape.dims(), |
| 82 | errors::InvalidArgument("axes ", axes_shape.DebugString(), |
| 83 | " can not have more elements" |
| 84 | " than input tensor has dimensions ", |
| 85 | x_shape.DebugString(), ".")); |
| 86 | // Reverse is a no-op if axes argument is empty. |
| 87 | if (axes_shape.num_elements() == 0) { |
| 88 | ctx->SetOutput(0, ctx->Input(0)); |
| 89 | return; |
| 90 | } |
| 91 | // XlaBuilder::Rev() requires concrete values for dimensions arg. |
| 92 | std::vector<int64> axes; |
| 93 | OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &axes)); |
| 94 | |
| 95 | // witnessed_axes is used to ensure that the same axis is not marked to be |
| 96 | // reversed multiple times. |
| 97 | absl::InlinedVector<bool, 8> witnessed_axes(x_shape.dims(), false); |
| 98 | |
| 99 | for (int d = 0; d < axes.size(); ++d) { |
| 100 | OP_REQUIRES( |
| 101 | ctx, (-x_shape.dims() <= axes[d]) && (axes[d] < x_shape.dims()), |
| 102 | errors::InvalidArgument(axes[d], " is out of range [-", |
| 103 | x_shape.dims(), ", ", x_shape.dims(), ").")); |
| 104 | // Axes can be negative and are shifted to the canonical index before |
| 105 | // being lowered to HLO. |
| 106 | if (axes[d] < 0) { |
| 107 | axes[d] += x_shape.dims(); |
| 108 | } |
| 109 | OP_REQUIRES(ctx, !witnessed_axes[axes[d]], |
| 110 | errors::InvalidArgument("canonicalized axis ", axes[d], |
| 111 | " was repeated.")); |
| 112 | witnessed_axes[axes[d]] = true; |
| 113 | } |
| 114 | |
| 115 | ctx->SetOutput(0, xla::Rev(ctx->Input(0), axes)); |
| 116 | } |
| 117 | }; |
| 118 | |
| 119 | REGISTER_XLA_OP(Name("ReverseV2").CompileTimeConstantInput("axis"), |
nothing calls this directly
no test coverage detected