| 25 | explicit RollOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} |
| 26 | |
| 27 | void Compile(XlaOpKernelContext* ctx) override { |
| 28 | const TensorShape input_shape = ctx->InputShape(0); |
| 29 | xla::XlaOp shift = ctx->Input(1); |
| 30 | const TensorShape shift_shape = ctx->InputShape(1); |
| 31 | const TensorShape axis_shape = ctx->InputShape(2); |
| 32 | |
| 33 | OP_REQUIRES(ctx, input_shape.dims() >= 1, |
| 34 | errors::InvalidArgument("input must be 1-D or higher")); |
| 35 | OP_REQUIRES(ctx, shift_shape.dims() <= 1, |
| 36 | errors::InvalidArgument( |
| 37 | "shift must be a scalar or a 1-D vector. Found: ", |
| 38 | shift_shape.DebugString())); |
| 39 | OP_REQUIRES( |
| 40 | ctx, shift_shape.dims() == axis_shape.dims(), |
| 41 | errors::InvalidArgument("shift and axis must have the same size")); |
| 42 | |
| 43 | xla::Literal axis; |
| 44 | OP_REQUIRES_OK(ctx, ctx->ConstantInput(2, &axis)); |
| 45 | |
| 46 | xla::XlaOp output = ctx->Input(0); |
| 47 | xla::PrimitiveType shift_type = ctx->input_xla_type(1); |
| 48 | int64 num_axes = axis_shape.dims() == 0 ? 1 : axis_shape.dim_size(0); |
| 49 | for (int64 i = 0; i != num_axes; ++i) { |
| 50 | int64 cur_axis = axis_shape.dims() == 0 ? *axis.GetIntegralAsS64({}) |
| 51 | : *axis.GetIntegralAsS64({i}); |
| 52 | |
| 53 | xla::XlaOp offset = |
| 54 | shift_shape.dims() == 0 |
| 55 | ? shift |
| 56 | : xla::Reshape(xla::SliceInDim(shift, /*start_index=*/i, |
| 57 | /*limit_index=*/i + 1, |
| 58 | /*stride=*/1, /*dimno=*/0), |
| 59 | {}); |
| 60 | xla::XlaOp axis_size = xla::ConstantR0WithType( |
| 61 | ctx->builder(), shift_type, input_shape.dim_size(cur_axis)); |
| 62 | // Adjust large offsets into [0, axis_size). This also makes negative |
| 63 | // offsets positive. |
| 64 | offset = ((offset % axis_size) + axis_size) % axis_size; |
| 65 | |
| 66 | // Stack two copies of the dimension, then slice from the calculated |
| 67 | // offset. |
| 68 | xla::XlaOp concat = |
| 69 | xla::ConcatInDim(ctx->builder(), {output, output}, cur_axis); |
| 70 | std::vector<xla::XlaOp> start_indices( |
| 71 | input_shape.dims(), xla::Zero(ctx->builder(), shift_type)); |
| 72 | start_indices[cur_axis] = axis_size - offset; |
| 73 | output = |
| 74 | xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes()); |
| 75 | } |
| 76 | ctx->SetOutput(0, output); |
| 77 | } |
| 78 | |
| 79 | private: |
| 80 | TF_DISALLOW_COPY_AND_ASSIGN(RollOp); |
nothing calls this directly
no test coverage detected