Function GetFlatShape

oneflow/user/kernels/gather_kernel.cpp:27–29

namespace {

Shape GetFlatShape(ShapeView shape, int64_t axis) {
  return Shape({shape.Count(0, axis), shape.At(axis), shape.Count(axis + 1)});
}

class GatherOpKernelCache final : public user_op::OpKernelCache {
 public: