Method InitOpKernelCache

oneflow/user/kernels/gather_kernel.cpp:67–83

  ~GatherKernel() override = default;

  std::shared_ptr<user_op::OpKernelCache> InitOpKernelCache(
      user_op::KernelCacheContext* ctx) const override {
    if (ctx->parallel_ctx().parallel_num() > 1) {
      const auto axis = ctx->Attr<int64_t>("axis");
      const NdSbp& in_nd_sbp = ctx->NdSbp4ArgNameAndIndex("in", 0);
      const Shape& hierarchy = *ctx->parallel_desc().hierarchy();
      CheckNdSbp(hierarchy, axis, in_nd_sbp, ctx->NdSbp4ArgNameAndIndex("indices", 0),
                 ctx->NdSbp4ArgNameAndIndex("out", 0));
      const Shape in_logical_shape =
          ExpandDimIf0D(ctx->LogicalTensorDesc4ArgNameAndIndex("in", 0)->shape());
      TensorSliceView view = GetTensorSliceView4ParallelId(hierarchy, in_nd_sbp, in_logical_shape,
                                                           ctx->parallel_ctx().parallel_id());
      return std::make_shared<GatherOpKernelCache>(view.At(axis).begin(), view.At(axis).end());
    } else {
      return nullptr;
    }
  }

 private:
  void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*,