Method Compile

tensorflow/compiler/tf2xla/kernels/reverse_op.cc:35–64

  explicit ReverseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}

  void Compile(XlaOpKernelContext* ctx) override {
    // r = tf.reverse(x, revdims)
    const TensorShape x_shape = ctx->InputShape(0);
    const TensorShape revd_shape = ctx->InputShape(1);
    // Validate input sizes.
    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(revd_shape),
                errors::InvalidArgument("axes must be a vector, not shape ",
                                        revd_shape.DebugString()));
    OP_REQUIRES(ctx, revd_shape.num_elements() == x_shape.dims(),
                errors::InvalidArgument("axes ", revd_shape.DebugString(),
                                        " must have same number of elements as"
                                        " than input tensor has dimensions ",
                                        x_shape.DebugString(), "."));
    if (revd_shape.num_elements() == 0) {
      ctx->SetOutput(0, ctx->Input(0));
      return;
    }
    // XlaBuilder::Rev() requires concrete values for dimensions arg.
    xla::Literal lax;
    OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &lax));

    std::vector<int64> dimensions;
    for (int d = 0; d < x_shape.dims(); ++d) {
      if (lax.Get<bool>({d})) {
        dimensions.push_back(d);
      }
    }

    ctx->SetOutput(0, xla::Rev(ctx->Input(0), dimensions));
  }
};

REGISTER_XLA_OP(Name("Reverse").CompileTimeConstantInput("dims"), ReverseOp);