Dev add contiguous view ops

oneflow/core/framework/tensor_methods.cpp

@@ -64,7 +64,6 @@ Maybe<Tensor> BasicView(const std::shared_ptr<Tensor>& input, const Shape& targe
 
 Maybe<Tensor> BasicView(const std::shared_ptr<Tensor>& input, const Shape& target_shape,
                         const Stride& target_stride, int64_t storage_offset) {
-  storage_offset = storage_offset + JUST(JUST(input->AsMirroredTensor())->storage_offset());
   // TODO(): Check shape compatible.
   auto device = JUST(input->device());
   auto tensor_meta = std::make_shared<MirroredTensorMeta>(
@@ -86,38 +85,14 @@ Maybe<Tensor> BasicView(const std::shared_ptr<Tensor>& input, const Shape& targe
   return output;
 }
 
-Maybe<Tensor> Reshape(const std::shared_ptr<Tensor>& input, const Shape& shape) {
+Maybe<Tensor> Reshape(const std::shared_ptr<Tensor>& input, const Shape& target_shape) {
   if (!(input->is_eager() && input->is_local())) {
     return Error::RuntimeError() << "view::Reshape(): input should be eager local tensor, but got "
                                  << (input->is_lazy() ? "lazy" : "consistent");
   }
-  int need_infer_axis = -1;
-  size_t count = 1;
-  for (int i = 0; i < shape.NumAxes(); ++i) {
-    if (shape.At(i) < -1) {
-      return Error::RuntimeError() << "Invalid shape dimension " << shape.At(i);
-    } else if (shape.At(i) == -1) {
-      CHECK_EQ_OR_RETURN(need_infer_axis, -1)
-          << "Shape " << shape.ToString() << " has more than 1 axis that needs to be infered.";
-      need_infer_axis = i;
-    } else {
-      count *= shape.At(i);
-    }
-  }
 
-  std::shared_ptr<Tensor> output;
-  size_t x_count = input->shape()->Count(0);
-  if (need_infer_axis == -1) {
-    CHECK_EQ_OR_RETURN(shape.Count(0), x_count);
-    output = JUST(BasicView(input, shape, 0));
-  } else {
-    Shape infered_shape = shape;
-    infered_shape.Set(need_infer_axis, x_count / count);
-    CHECK_EQ_OR_RETURN(infered_shape.Count(0), x_count)
-        << "Shape " << shape.ToString() << " is invalid for input of shape "
-        << input->shape()->ToString();
-    output = JUST(BasicView(input, infered_shape, 0));
-  }
+  int64_t storage_offset = JUST(JUST(input->AsMirroredTensor())->storage_offset());
+  std::shared_ptr<Tensor> output = JUST(BasicView(input, target_shape, storage_offset));
 
   if (autograd::GradMode::is_enabled() && input->requires_grad()) {
     Shape input_shape(input->shape()->dim_vec());
@@ -192,6 +167,100 @@ Maybe<Tensor> Slice(const std::shared_ptr<Tensor>& input, const std::vector<int6
   return output;
 }
 
+
+Maybe<Tensor> UnSqueeze(const std::shared_ptr<Tensor>& input, const int32_t& expand_dim) {
+  if (!(input->is_eager() && input->is_local())) {
+    return Error::RuntimeError()
+           << "view::UnSqueeze(): input should be eager local tensor, but got "
+           << (input->is_lazy() ? "lazy" : "consistent");
+  }
+
+  const auto& shape = input->shape();
+  const auto& strides = JUST(input->stride());
+  const auto& ndim = shape->NumAxes();
+
+  DimVector target_dim_vec(ndim + 1);
+  StrideVector target_stride_vec(ndim + 1);
+
+  int cnt = 0;
+  for (int i = 0; i < ndim; i++) {
+    if (i == expand_dim) { cnt++; }
+    target_dim_vec[cnt] = shape->At(i);
+    target_stride_vec[cnt] = strides->At(i);
+    cnt++;
+  }
+  target_dim_vec[expand_dim] = 1;
+  target_stride_vec[expand_dim] = strides->At(expand_dim);
+
+  int64_t storage_offset = JUST(JUST(input->AsMirroredTensor())->storage_offset());
+  std::shared_ptr<Tensor> output =
+      JUST(BasicView(input, Shape(target_dim_vec), Stride(target_stride_vec), storage_offset));
+
+  if (autograd::GradMode::is_enabled() && input->requires_grad()) {
+    auto backward_fn =
+        std::make_shared<std::function<Maybe<void>(const TensorTuple&, TensorTuple*, bool)>>(
+            [=](const TensorTuple& out_grads, TensorTuple* in_grads,
+                bool create_graph) -> Maybe<void> {
+              autograd::AutoGradMode mode(create_graph);
+              CHECK_EQ_OR_RETURN(out_grads.size(), 1);
+              in_grads->resize(1);
+              in_grads->at(0) = JUST(functional::Reshape(out_grads.at(0), *shape));
+              return Maybe<void>::Ok();
+            });
+    TensorTuple outputs{output};
+    JUST(GetThreadLocalAutogradEngine()->AddBackwardFuncPtr("view::unsqueeze_backward",
+                                                            backward_fn, {input}, &outputs));
+  }
+  return output;
+}
+
+Maybe<Tensor> Squeeze(const std::shared_ptr<Tensor>& input,
+                      const std::vector<int32_t>& squeeze_dims) {
+  if (!(input->is_eager() && input->is_local())) {
+    return Error::RuntimeError() << "view::Squeeze(): input should be eager local tensor, but got "
+                                 << (input->is_lazy() ? "lazy" : "consistent");
+  }
+
+  const auto& shape = input->shape();
+  const auto& strides = JUST(input->stride());
+  const int64_t ndim = shape->NumAxes();
+
+  const int target_ndim = ndim - squeeze_dims.size();
+  DimVector target_dim_vec(target_ndim);
+  StrideVector target_stride_vec(target_ndim);
+
+  int cnt = 0;
+  for (int i = 0; i < ndim; i++) {
+    if (find(squeeze_dims.begin(), squeeze_dims.end(), i) == squeeze_dims.end()) {
+      target_dim_vec[cnt] = shape->At(i);
+      target_stride_vec[cnt] = strides->At(i);
+      cnt++;
+    }
+  }
+
+  int64_t storage_offset = JUST(JUST(input->AsMirroredTensor())->storage_offset());
+  std::shared_ptr<Tensor> output =
+      JUST(BasicView(input, Shape(target_dim_vec), Stride(target_stride_vec), storage_offset));
+
+  if (autograd::GradMode::is_enabled() && input->requires_grad()) {
+    auto backward_fn =
+        std::make_shared<std::function<Maybe<void>(const TensorTuple&, TensorTuple*, bool)>>(
+            [=](const TensorTuple& out_grads, TensorTuple* in_grads,
+                bool create_graph) -> Maybe<void> {
+              autograd::AutoGradMode mode(create_graph);
+              CHECK_EQ_OR_RETURN(out_grads.size(), 1);
+              in_grads->resize(1);
+              in_grads->at(0) = JUST(functional::ReshapeLike(out_grads.at(0), input));
+              return Maybe<void>::Ok();
+            });
+    TensorTuple outputs{output};
+    JUST(GetThreadLocalAutogradEngine()->AddBackwardFuncPtr("view::squeeze_backward", backward_fn,
+                                                            {input}, &outputs));
+  }
+  return output;
+}
+
+
 }  // namespace view
 }  // namespace one
 }  // namespace oneflow

oneflow/core/framework/tensor_methods.h

@@ -29,13 +29,21 @@ Maybe<bool> IsContiguous(const std::shared_ptr<Tensor>& tensor);
 namespace view {
 
 Maybe<Tensor> BasicView(const std::shared_ptr<Tensor>& input, const Shape& target_shape,
-                        const Stride& target_strides, int64_t storage_offset);
+                        int64_t storage_offset);
+
+Maybe<Tensor> BasicView(const std::shared_ptr<Tensor>& input, const Shape& target_shape,
+                        const Stride& target_stride, int64_t storage_offset);
 
 Maybe<Tensor> Reshape(const std::shared_ptr<Tensor>& input, const Shape& shape);
 
 Maybe<Tensor> Slice(const std::shared_ptr<Tensor>& input, const std::vector<int64_t>& starts,
                     const std::vector<int64_t>& ends, const std::vector<int64_t>& steps);
 
+Maybe<Tensor> UnSqueeze(const std::shared_ptr<Tensor>& input, const int32_t& expand_dim);
+
+Maybe<Tensor> Squeeze(const std::shared_ptr<Tensor>& input,
+                      const std::vector<int32_t>& squeeze_dims);
+
 }  // namespace view
 }  // namespace one
 }  // namespace oneflow

oneflow/core/functional/functional_api.yaml

@@ -784,6 +784,12 @@
   signature: "Tensor (Tensor input, Int32 dim) => Unsqueeze"
   bind_python: True
 
+- name: "squeeze"
+  signature: [
+    "Tensor (Tensor x, Int32List[1] dim=None) => Squeeze",
+  ]
+  bind_python: True
+
 - name: "exp"
   signature: "Tensor (Tensor x) => Exp"
   bind_python: True
@@ -1107,6 +1113,10 @@
   signature: "Tensor (Tensor x, Shape shape) => Reshape"
   bind_python: True
 
+- name: "view"
+  signature: "Tensor (Tensor x, Shape shape) => View"
+  bind_python: True
+
 - name: "slice_view_1d_contiguous"
   signature: "Tensor (Tensor x, Int64 start, Int64 end) => SliceView1dContiguous"
   bind_python: True
@@ -1139,12 +1149,6 @@
   signature: "Void (Tensor ref, Tensor value, Int64List start, Int64List stop, Int64List step) => LogicalSliceAssign"
   bind_python: True
 
-- name: "squeeze"
-  signature: [
-    "Tensor (Tensor x, Int32List[1] dim=None) => Squeeze",
-  ]
-  bind_python: True
-
 - name: "copy"
   signature: "Tensor (Tensor x, String device_type, Int64 device_id) => Copy"
   bind_python: True