Method Compile

tensorflow/compiler/tf2xla/kernels/sort_ops.cc:28–33

  explicit XlaSortOp(OpKernelConstruction* context) : XlaOpKernel(context) {}

  void Compile(XlaOpKernelContext* context) override {
    context->SetOutput(0, xla::Sort({context->Input("input")},
                                    xla::CreateScalarLtComputation(
                                        {context->InputXlaType("input")},
                                        context->builder())));
  }
};

REGISTER_XLA_OP(Name("XlaSort"), XlaSortOp);