symbolic shape

oneflow/api/cpp/framework/graph.cpp

@@ -332,8 +332,12 @@ of::Maybe<void> Graph::GraphImpl::AddOp(of::OperatorConf op_conf) {
   }
   op_conf.set_device_tag(GetDeviceTag(device_));
   if (batch_size_ > 0 && op_conf.has_input_conf()) {
-    op_conf.mutable_input_conf()->mutable_blob_conf()->mutable_shape()->mutable_dim()->Set(
-        0, batch_size_);
+    op_conf.mutable_input_conf()
+        ->mutable_blob_conf()
+        ->mutable_shape()
+        ->mutable_dim()
+        ->Add()
+        ->set_int64_value(batch_size_);
   }
   auto* ctx = JUST(of::GetCurInferCtx());
   JUST(ctx->AddAndInferGlobalOp(op_conf));
@@ -379,7 +383,7 @@ of::Maybe<void> Graph::GraphImpl::BuildGraph() {
           const std::string input_lbi_str = op_conf.output_conf().in();
           const of::LogicalBlobId input_lbi = of::GenLogicalBlobId(input_lbi_str);
           int64_t batch_size = node->LogicalBlobDesc4Lbi(input_lbi).shape().At(0);
-          blob_conf.mutable_shape()->set_dim(0, batch_size);
+          blob_conf.mutable_shape()->mutable_dim()->Add()->set_int64_value(batch_size);
         }
         output_name_to_tensor_[op_conf.name()] = JUST(of::one::functional::Empty(
             of::Shape(blob_conf.shape()),

oneflow/api/python/framework/dim.cpp

@@ -0,0 +1,65 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include <pybind11/pybind11.h>
+#include <pybind11/operators.h>
+#include "oneflow/api/python/of_api_registry.h"
+#include "oneflow/core/common/dim.h"
+
+namespace py = pybind11;
+
+namespace oneflow {
+
+#define DEFINE_OPERATOR(op)                           \
+  .def(py::self op py::self)                          \
+      .def(py::self op char())                        \
+      .def(py::self op static_cast<unsigned char>(0)) \
+      .def(py::self op int())                         \
+      .def(py::self op static_cast<unsigned int>(0))  \
+      .def(py::self op long())                        \
+      .def(py::self op 0UL)                           \
+      .def(py::self op 0LL)                           \
+      .def(py::self op 0ULL)                          \
+      .def(py::self op float())                       \
+      .def(py::self op double())
+
+ONEFLOW_API_PYBIND11_MODULE("", m) {
+  py::class_<Dim>(m, "Dim")
+      .def(py::init([](int value) { return Dim(value); }))
+      .def_static("unknown", &Dim::Unknown)
+      .def("__str__",
+           [](Dim dim) {
+             std::stringstream ss;
+             ss << dim;
+             return ss.str();
+           })
+      .def("__repr__",
+           [](Dim dim) {
+             std::stringstream ss;
+             ss << dim;
+             return ss.str();
+           })
+      // clang-format off
+      DEFINE_OPERATOR(==)
+      DEFINE_OPERATOR(!=)
+      DEFINE_OPERATOR(<)
+      DEFINE_OPERATOR(<=)
+      DEFINE_OPERATOR(>)
+      DEFINE_OPERATOR(>=)
+      // clang-format on
+      .def(py::hash(py::self));
+}
+
+}  // namespace oneflow

oneflow/api/python/framework/size.cpp

@@ -31,7 +31,7 @@ static PyObject* TensorSize_repr(TensorSize* self) {
   int32_t size = PyTuple_Size((PyObject*)self);
   ss << "oneflow.Size([";
   for (int i = 0; i < size; ++i) {
-    int64_t dim = PyLong_AsLongLong(PyTuple_GET_ITEM(self, i));
+    Dim dim = one::functional::PyUnpackShapeItem(PyTuple_GET_ITEM(self, i));
     ss << dim;
     if (++idx != size) { ss << ", "; }
   }
@@ -44,10 +44,11 @@ static PyObject* TensorSize_new(PyTypeObject* type, PyObject* args, PyObject* kw
   if (self.get()) {
     for (int i = 0; i < PyTuple_Size(self.get()); ++i) {
       PyObject* item = PyTuple_GET_ITEM(self.get(), i);
-      if (!PyLong_Check(item)) {
-        return PyErr_Format(PyExc_TypeError,
-                            "oneflow.Size() takes an iterable of 'int', but item '%d' is '%s'", i,
-                            Py_TYPE(item)->tp_name);
+      if (!one::functional::PyShapeItemCheck(item)) {
+        return PyErr_Format(
+            PyExc_TypeError,
+            "oneflow.Size() takes an iterable of 'int' or 'Dim', but item '%d' is '%s'", i,
+            Py_TYPE(item)->tp_name);
       }
     }
   }
@@ -121,8 +122,18 @@ static PyObject* TensorSize_numel(PyObject* self, PyObject* args) {
   return PyLong_FromLongLong(numel);
 }
 
+static PyObject* TensorSize_all_dims_known(PyObject* self, PyObject* args) {
+  for (int i = 0; i < PyTuple_Size(self); ++i) {
+    Dim dim = one::functional::PyUnpackShapeItem(PyTuple_GET_ITEM((TensorSize*)self, i));
+    if (!dim.is_known()) { Py_RETURN_FALSE; }
+  }
+  Py_RETURN_TRUE;
+}
+
 static PyMethodDef TensorSize_methods[] = {
-    {"numel", (PyCFunction)TensorSize_numel, METH_NOARGS, NULL}, {NULL}};
+    {"numel", (PyCFunction)TensorSize_numel, METH_NOARGS, NULL},
+    {"all_dims_known", (PyCFunction)TensorSize_all_dims_known, METH_NOARGS, NULL},
+    {NULL}};
 
 PyTypeObject TensorSize_Type = {
     PyVarObject_HEAD_INIT(NULL, 0) "oneflow.Size", /* tp_name */
@@ -173,7 +184,10 @@ PyObject* TensorSize_NewFromShape(const Shape& size) {
   PyObjectPtr self(TensorSize_New(size.NumAxes()));
   if (self.get()) {
     for (int i = 0; i < size.NumAxes(); ++i) {
-      PyTuple_SET_ITEM(self.get(), i, PyLong_FromLongLong(size.At(i)));
+      Dim dim = size[i];
+      PyTuple_SET_ITEM(self.get(), i,
+                       dim.is_known() ? one::functional::CastToPyObject(dim.val())
+                                      : one::functional::CastToPyObject(dim));
     }
   }
   return self.release();
@@ -186,11 +200,11 @@ Shape TensorSize_AsShape(PyObject* self) {
     return Shape();
   }
   int size = TensorSize_length((TensorSize*)self);
-  DimVector dim_vec(size);
+  Shape shape(size);
   for (int i = 0; i < size; ++i) {
-    dim_vec[i] = PyLong_AsLongLong(PyTuple_GET_ITEM((TensorSize*)self, i));
+    shape[i] = py::handle(PyTuple_GET_ITEM((TensorSize*)self, i)).cast<Dim>();
   }
-  return Shape(std::move(dim_vec));
+  return shape;
 }
 
 ONEFLOW_API_PYBIND11_MODULE("", m) {

oneflow/api/python/functional/common.cpp

@@ -248,25 +248,50 @@ std::vector<Symbol<DType>> PyUnpackDTypeSequence(PyObject* obj) {
   return PyUnpackSequence<Symbol<DType>>(obj, [](PyObject* item) { return PyUnpackDType(item); });
 }
 
+bool PyDimCheck(PyObject* obj) { return detail::isinstance_fast<Dim>(obj); }
+Dim PyUnpackDim(PyObject* obj) { return *detail::cast_fast<Dim*>(obj); }
+
 // Shape
-bool PyShapeCheck(PyObject* obj) { return PyLongSequenceCheck(obj); }
+bool PyShapeItemCheck(PyObject* obj) {
+  return PyLong_Check(obj) || PyIntegerScalarTensorCheck(obj) || PyDimCheck(obj);
+}
+Dim PyUnpackShapeItem(PyObject* obj) {
+  if (PyLong_Check(obj)) {
+    return PyLong_AsLongLong(obj);
+  } else if (PyDimCheck(obj)) {
+    return PyUnpackDim(obj);
+  } else {
+    return PyUnpackIntegerScalarTensor_AsLongLong(obj);
+  }
+}
+bool PyShapeCheck(PyObject* obj) {
+  return PySequenceCheck(obj, [](PyObject* item) { return PyShapeItemCheck(item); });
+}
 
 Shape PyUnpackShape(PyObject* obj) {
   bool is_tuple = PyTuple_Check(obj);
   CHECK_OR_THROW(is_tuple || PyList_Check(obj))
       << "The object is not list or tuple, but is " << Py_TYPE(obj)->tp_name;
   size_t size = is_tuple ? PyTuple_GET_SIZE(obj) : PyList_GET_SIZE(obj);
-  DimVector values(size);
+  Shape shape(size);
   for (int i = 0; i < size; ++i) {
     PyObject* item = is_tuple ? PyTuple_GET_ITEM(obj, i) : PyList_GET_ITEM(obj, i);
-    values[i] = PyLong_AsLongLong(item);
+    if (PyLong_Check(item)) {
+      shape.Set(i, PyLong_AsLongLong(item));
+    } else if (PyIntegerScalarTensorCheck(item)) {
+      shape.Set(i, PyUnpackIntegerScalarTensor_AsLongLong(item));
+    } else if (PyDimCheck(item)) {
+      shape.Set(i, PyUnpackDim(item));
+    } else {
+      THROW(RuntimeError) << "Unsupported item type: " << Py_TYPE(item)->tp_name;
+    }
   }
-  return Shape(values);
+  return shape;
 }
 
 // Shape list
 bool PyShapeSequenceCheck(PyObject* obj) {
-  return PySequenceCheck(obj, [](PyObject* item) { return PyLongSequenceCheck(item); });
+  return PySequenceCheck(obj, [](PyObject* item) { return PyShapeCheck(item); });
 }
 std::vector<Shape> PyUnpackShapeSequence(PyObject* obj) {
   return PyUnpackSequence<Shape>(obj, [](PyObject* item) -> Shape { return PyUnpackShape(item); });

oneflow/api/python/functional/common.h

@@ -22,8 +22,7 @@ limitations under the License.
 #include <pybind11/pybind11.h>
 
 #include "oneflow/api/python/framework/tensor.h"
-#include "oneflow/api/python/caster/maybe.h"
-#include "oneflow/api/python/caster/optional.h"
+#include "oneflow/api/python/of_api_registry.h"
 #include "oneflow/core/common/throw.h"
 #include "oneflow/core/common/maybe.h"
 #include "oneflow/core/common/preprocessor.h"
@@ -166,7 +165,13 @@ Symbol<MemoryFormat> PyUnpackMemoryFormat(PyObject* obj);
 bool PyDTypeSequenceCheck(PyObject* obj);
 std::vector<Symbol<DType>> PyUnpackDTypeSequence(PyObject* obj);
 
+// Dim
+bool PyDimCheck(PyObject* obj);
+Dim PyUnpackDim(PyObject* obj);
+
 // Shape
+bool PyShapeItemCheck(PyObject* obj);
+Dim PyUnpackShapeItem(PyObject* obj);
 bool PyShapeCheck(PyObject* obj);
 Shape PyUnpackShape(PyObject* obj);
 

oneflow/api/python/functional/dispatch_stateful_ops.cpp

@@ -32,13 +32,16 @@ namespace functional {
 namespace impl {
 
 ONEFLOW_FUNCTION_LIBRARY(m) {
-  m.add_functor(
-      "DispatchFeedInput",
-      [](const std::shared_ptr<OpExpr>& op, const std::shared_ptr<Tensor>& input) -> Maybe<Tensor> {
-        const auto& origin_input = JUST(OpInterpUtil::Dispatch<Tensor>(*op, {input}));
-        // Unpack input when do grad acc
-        return GradAccTryInsertUnpackAfterInput(origin_input);
-      });
+  m.add_functor("DispatchFeedInput",
+                [](const std::shared_ptr<OpExpr>& op, const std::shared_ptr<Tensor>& input,
+                   const Optional<Shape>& shape) -> Maybe<Tensor> {
+                  const auto& origin_input =
+                      shape.has_value() ? JUST(OpInterpUtil::Dispatch<Tensor>(
+                          *op, {input}, OpExprInterpContext(JUST(shape))))
+                                        : JUST(OpInterpUtil::Dispatch<Tensor>(*op, {input}));
+                  // Unpack input when do grad acc
+                  return GradAccTryInsertUnpackAfterInput(origin_input);
+                });
   m.add_functor(
       "DispatchFetchOutput",
       [](const std::shared_ptr<OpExpr>& op, const std::shared_ptr<Tensor>& input) -> Maybe<Tensor> {

oneflow/api/python/functional/dispatch_stateful_ops.yaml

@@ -21,7 +21,7 @@
 # }
 
 - name: "dispatch_feed_input"
-  signature: "Tensor (OpExpr op, Tensor input) => DispatchFeedInput"
+  signature: "Tensor (OpExpr op, Tensor input, Shape shape=None) => DispatchFeedInput"
   bind_python: True
 
 - name: "dispatch_feed_variable"

oneflow/api/python/functional/python_arg.cpp

@@ -243,7 +243,7 @@ bool PythonArg::TypeCheck(ValueType type) const {
     case kDTYPE: return PyDTypeCheck(object_);
     case kLAYOUT: return PyLayoutCheck(object_);
     case kMEMORYFORMAT: return PyMemoryFormatCheck(object_);
-    case kSHAPE: return PyLongSequenceCheck(object_);
+    case kSHAPE: return PyShapeCheck(object_);
     case kGENERATOR:
     case kGENERATOR_REF: return PyGeneratorCheck(object_);
     case kTENSOR_INDEX: return PyTensorIndexCheck(object_);

oneflow/core/autograd/gradient_funcs/expand.cpp

@@ -57,7 +57,7 @@ Maybe<void> Expand::Capture(ExpandCaptureState* ctx, const TensorTuple& inputs,
   if (ctx->keep_dims) {
     for (size_t i = 0; i < expand_shape.size(); ++i) {
       const auto& t_dim = expand_shape[i];
-      const auto& dim = i < ctx->lpad ? 1 : in_shape[i - ctx->lpad];
+      const auto& dim = i < ctx->lpad ? 1 : in_shape[i - ctx->lpad].val();
       if (dim != t_dim) { ctx->reduce_dims.push_back(i); }
     }
   } else {

oneflow/core/autograd/gradient_funcs/matmul.cpp

@@ -153,7 +153,7 @@ Maybe<void> BroadcastMatmul::Capture(BroadcastMatmulCaptureState* ctx, const Ten
     const int64_t num_batch_dims = num_dims - 2;
     const int64_t num_padding_dims = num_max_batch_dims - num_batch_dims;
     return [num_padding_dims, shape_dim](size_t index) {
-      return index < num_padding_dims ? 1 : shape_dim.At(index - num_padding_dims);
+      return index < num_padding_dims ? Dim(1) : shape_dim.At(index - num_padding_dims);
     };
   };
   auto GetABatchDim = MakeGetBatchDim(a_num_axes, *a_shape);
@@ -218,7 +218,7 @@ Maybe<void> BroadcastMatmul::Apply(const BroadcastMatmulCaptureState* ctx,
     const int64_t num_batch_dims = num_dims - 2;
     const int64_t num_padding_dims = num_max_batch_dims - num_batch_dims;
     return [num_padding_dims, shape_dim](size_t index) {
-      return index < num_padding_dims ? 1 : shape_dim.At(index - num_padding_dims);
+      return index < num_padding_dims ? Dim(1) : shape_dim.At(index - num_padding_dims);
     };
   };
   auto GetOutBatchDim = MakeGetBatchDim(out_num_axes, *out_shape);

oneflow/core/boxing/boxing_dividor_util.cpp

@@ -244,7 +244,7 @@ Maybe<Symbol<ParallelDesc>> GetFisrtDeviceOfPlacement(Symbol<ParallelDesc> place
   parallel_conf.add_device_name(std::string("@") + std::to_string(machine_id) + ":"
                                 + std::to_string(device_id));
   for (int64_t i = 0; i < placement->hierarchy()->NumAxes(); ++i) {
-    parallel_conf.mutable_hierarchy()->add_dim(1);
+    parallel_conf.mutable_hierarchy()->add_dim()->set_int64_value(1);
   }
   std::shared_ptr<ParallelDesc> parallel_desc;
   JUST(PhysicalRun([&parallel_desc, &parallel_conf](InstructionsBuilder* builder) -> Maybe<void> {