Develop_OneFlow_Normal_Distribution_Synced_With_PyTorch_Tensor

oneflow/core/functional/functional_api.yaml

@@ -839,8 +839,18 @@
   signature: "Tensor (Tensor x, Float lower=0.125, Float upper=0.3333333333333333, Bool training=False) => RReluInplace"
   bind_python: True
 
+- name: "normal_"
+  signature: "Tensor (Tensor x, Float mean=0.0, Float std=1.0, Generator generator=None) => Normal_"
+  bind_python: True
+
 - name: "normal"
   signature: [
+      "Tensor (Tensor mean, Tensor std, *, Tensor out=None, 
+      Generator generator=None, Bool requires_grad=False) => TensorTensorNormal",
+      "Tensor (Tensor mean, Float std=1.0, *, Tensor out=None, 
+      Generator generator=None, Bool requires_grad=False) => TensorScalarNormal",
+      "Tensor (Float mean, Tensor std, *, Tensor out=None, 
+      Generator generator=None, Bool requires_grad=False) => ScalarTensorNormal",
       "Tensor (Float mean, Float std, Shape size, *, Tensor out=None, DataType dtype=None, Device device=None,
       Generator generator=None, Bool requires_grad=False) => Normal",
       "Tensor (Float mean, Float std, Int32 size, *, Tensor out=None, DataType dtype=None, Device device=None,

oneflow/core/functional/impl/random_functor.cpp

@@ -13,14 +13,28 @@ 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 <google/protobuf/map_entry_lite.h>
+#include <cstdint>
+#include "oneflow/core/common/data_type.pb.h"
+#include "oneflow/core/common/just.h"
+#include "oneflow/core/common/maybe.h"
+#include "oneflow/core/common/scalar.h"
+#include "oneflow/core/common/shape.h"
+#include "oneflow/core/common/util.h"
+#include "oneflow/core/framework/device.h"
+#include "oneflow/core/framework/dtype.h"
 #include "oneflow/core/framework/layout.h"
 #include "oneflow/core/framework/mutable_attr_map.h"
 #include "oneflow/core/framework/nd_sbp.h"
 #include "oneflow/core/framework/op_builder.h"
+#include "oneflow/core/framework/tensor.h"
+#include "oneflow/core/framework/tensor_util.h"
 #include "oneflow/core/functional/function_library.h"
+#include "oneflow/core/functional/functional_api.yaml.h"
 #include "oneflow/core/functional/impl/unary_functor.h"
 #include "oneflow/core/job/global_mode.h"
 #include "oneflow/core/job/parallel_desc.h"
+#include "oneflow/core/kernel/kernel_util.h"
 #include "oneflow/user/kernels/distributions/common.h"
 #include "oneflow/user/kernels/random_seed_util.h"
 #include "oneflow/core/rpc/include/global_process_ctx.h"
@@ -316,6 +330,10 @@ class NormalFunctor {
 
     if (out.has_value()) {
       auto out_tensor = JUST(out);
+
+      CHECK_OR_RETURN(shape == (*out_tensor->shape()))
+          << Error::RuntimeError() << "Shape of out_tensor does not match shape";
+
       Symbol<DType> output_tensor_dtype = out_tensor->dtype();
       if (optional_dtype.has_value()) {
         CHECK_OR_RETURN(output_tensor_dtype == dtype)
@@ -371,6 +389,117 @@ class Normal2Functor {
   }
 };
 
+class InplaceNormalFuctor {
+ public:
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const float& mean,
+                           const float& std,
+                           const Optional<one::Generator>& optional_generator) const {
+    return Normal(mean, std, *x->shape(), x, x->dtype(), JUST(x->device()), optional_generator,
+                  x->requires_grad());
+  }
+};
+
+class NormalTools {
+ public:
+  template<typename T>
+  static T GetTensorItemValue(const std::shared_ptr<one::Tensor>& input) {
+    // CHECK_EQ(input->nelement(), 1) << "Input tensor must have exactly one element";
+    std::unique_ptr<T> ptr(new T());
+    const auto& callback = [&](ep::Stream* stream,
+                               const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object) {
+      SyncAutoMemcpy(stream, ptr.get(), eager_blob_object->dptr(), sizeof(T),
+                     memory::MakeHostMemCase(), eager_blob_object->mem_case());
+    };
+    SyncAccessTensorWithTimeOut(input, callback, "const").GetOrThrow();
+    return *ptr;
+  }
+
+  static Shape TensorBroadcastShape(const std::shared_ptr<one::Tensor>& mean,
+                                    const std::shared_ptr<one::Tensor>& std) {
+    int64_t dimsA = mean->ndim();
+    int64_t dimsB = std->ndim();
+    int64_t ndim = dimsA > dimsB ? dimsA : dimsB;
+    Shape out_shape(ndim);
+
+    // Use ptrdiff_t to ensure signed comparison.
+    for (ptrdiff_t i = (ptrdiff_t)ndim - 1; i >= 0; --i) {
+      ptrdiff_t offset = ndim - 1 - i;
+      ptrdiff_t dimA = dimsA - 1 - offset;
+      ptrdiff_t dimB = dimsB - 1 - offset;
+      int64_t sizeA = (dimA >= 0) ? mean->dim(dimA) : 1;
+      int64_t sizeB = (dimB >= 0) ? std->dim(dimB) : 1;
+
+      CHECK_OR_THROW(sizeA == sizeB || sizeA == 1 || sizeB == 1)
+          << "The size of tensor a (" << sizeA << ") must match the size of tensor b (" << sizeB
+          << ") at non-singleton dimension " << i;
+
+      // 1s map to the other size (even 0).
+      out_shape.Set(i, sizeA == 1 ? sizeB : sizeA);
+    }
+    return out_shape;
+  }
+
+  static bool CheckNormalTensorStd(const std::shared_ptr<one::Tensor>& std) {
+    if (std->nelement() == 0) { return true; }
+    GetTensorItemValue<float>(CHECK_JUST(Min(std)));
+    auto std_check = CHECK_JUST(ScalarLogicalGreaterEqual(CHECK_JUST(Min(std)), Scalar(0.0)));
+    GetTensorItemValue<bool>(std_check);
+    CHECK(GetTensorItemValue<bool>(std_check)) << "normal expects all elements of std >= 0.0";
+    return true;
+  }
+};
+
+class TensorTensorNormalFunctor {
+ public:
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& mean,
+                           const std::shared_ptr<one::Tensor>& std,
+                           const Optional<one::Tensor>& out,
+                           const Optional<one::Generator>& optional_generator,
+                           const bool& requires_grad) const {
+    auto out_shape = impl::NormalTools::TensorBroadcastShape(mean, std);
+    impl::NormalTools::CheckNormalTensorStd(std);  // NOLINT
+    auto output = JUST(Normal(0, 1, out_shape, out, Symbol<DType>(mean->dtype()),
+                              JUST(mean->device()), optional_generator, requires_grad));
+    // mean + output * std
+    JUST(InplaceMul(output, std));
+    JUST(Add(output, mean, 1, true));
+    JUST(output->set_requires_grad(requires_grad));
+    return output;
+  }
+};
+
+class TensorScalarNormalFunctor {
+ public:
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& mean, const float& std,
+                           const Optional<one::Tensor>& out,
+                           const Optional<one::Generator>& optional_generator,
+                           const bool& requires_grad) const {
+    CHECK_GE_OR_RETURN(std, 0.0) << "normal expects std >= 0.0, but found std " << (std)
+                                 << ". This may cause an error.";
+    auto output = JUST(Normal(0, std, *(mean->shape()), out, mean->dtype(), JUST(mean->device()),
+                              optional_generator, requires_grad));
+    JUST(Add(output, mean, 1, true));
+    JUST(output->set_requires_grad(requires_grad));
+    return output;
+  }
+};
+
+class ScalarTensorNormalFunctor {
+ public:
+  Maybe<Tensor> operator()(const float& mean, const std::shared_ptr<one::Tensor>& std,
+                           const Optional<one::Tensor>& out,
+                           const Optional<one::Generator>& optional_generator,
+                           const bool& requires_grad) const {
+    impl::NormalTools::CheckNormalTensorStd(std);  // NOLINT
+    auto output = JUST(Normal(0.0, 1.0, *(std->shape()), out, std->dtype(), JUST(std->device()),
+                              optional_generator, requires_grad));
+    JUST(InplaceMul(output, std));
+    JUST(ScalarAdd(output, mean, 1, true));
+    JUST(output->set_requires_grad(requires_grad));
+    return output;
+  }
+};
+
 class GlobalNormalFunctor {
  public:
   GlobalNormalFunctor() { op_ = CHECK_JUST(one::OpBuilder("normal").Output("out").Build()); }
@@ -835,6 +964,10 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<GlobalRandNFunctor>("GlobalRandN");
   m.add_functor<impl::NormalFunctor>("Normal");
   m.add_functor<impl::Normal2Functor>("Normal2");
+  m.add_functor<impl::TensorTensorNormalFunctor>("TensorTensorNormal");
+  m.add_functor<impl::TensorScalarNormalFunctor>("TensorScalarNormal");
+  m.add_functor<impl::ScalarTensorNormalFunctor>("ScalarTensorNormal");
+  m.add_functor<impl::InplaceNormalFuctor>("Normal_");
   m.add_functor<impl::GlobalNormalFunctor>("GlobalNormal");
   m.add_functor<impl::GlobalNormal2Functor>("GlobalNormal2");
   m.add_functor<RandnLikeFunctor>("RandnLike");

python/oneflow/__init__.py

@@ -428,6 +428,7 @@ def atexit_hook(hook):
 from oneflow.nn.modules.meshgrid import meshgrid_op as meshgrid
 from oneflow.nn.modules.unique import unique_op as unique
 from oneflow._C import normal
+from oneflow._C import normal_
 from oneflow._C import rand
 from oneflow._C import randn
 from oneflow._C import randn_like

python/oneflow/framework/docstr/random.py

@@ -193,8 +193,103 @@
 add_docstr(
     oneflow._C.normal,
     r"""
-    normal(mean, std, size, *, out=None, placement=None, sbp=None, generator=None, dtype=None, device=None, requires_grad=False) -> Tensor
+    normal(mean, std, *, generator=None, out=None) -> Tensor
+
+    Returns a tensor of random numbers drawn from separate normal distributions
+    whose mean and standard deviation are given.
+
+    The :attr:`mean` is a tensor with the mean of
+    each output element's normal distribution
+
+    The :attr:`std` is a tensor with the standard deviation of
+    each output element's normal distribution
+
+    The shapes of :attr:`mean` and :attr:`std` don't need to match, but the
+    total number of elements in each tensor need to be the same.
+
+    .. note:: Infers the output shape from input arrays :attr:`mean` and :attr:`std`.
+    The output shape will have a dimensionality equal to the max of :attr:`mean` and :attr:`std`.
+    Dimensions with size 1 in either :attr:`mean` or :attr:`std` are expanded to match the other.
+
+    Args:
+        mean (Tensor): the tensor of per-element means
+        std (Tensor): the tensor of per-element standard deviations
+
+    Keyword args:
+        generator (Generator, optional): Random number generator. Defaults to `oneflow::DefaultGenerator` if not provided.
+        out (Tensor, optional): Output tensor, will be resized and filled with the result. If not provided, a new tensor is created.  
+
+    Example::
+
+    .. code-block:: python
+
+        >>> import oneflow as flow
+        >>> generator = flow.Generator()
+        >>> generator.manual_seed(0) #doctest: +ELLIPSIS
+        <oneflow._oneflow_internal.Generator object at ...>
+        >>> z = flow.normal(mean=flow.arange(1., 11.), std=flow.arange(1, 0, -0.1), generator=generator)
+        >>> z[:5]
+        tensor([3.2122, 3.0468, 3.6192, 4.3387, 5.6261], dtype=oneflow.float32)
+
+    normal(mean: Tensor, std: double, *, generator: Optional[Generator] = None, out: Optional[Tensor] = None) -> Tensor  
+
+    Returns a tensor of random numbers drawn from separate normal distributions
+    whose mean and standard deviation are given.  
+
+    The :attr:`mean` is a tensor with the mean of  
+    each output element's normal distribution  
+
+    :attr:`std` is a double with the standard deviation of  
+    each output element's normal distribution  
+
+    Args:  
+        mean (Tensor): the tensor of per-element means  
+        std (double): the standard deviation  
 
+    Keyword args:  
+        generator (Generator, optional): Random number generator. Defaults to `oneflow::DefaultGenerator` if not provided.
+        out (Tensor, optional): Output tensor, will be resized and filled with the result. If not provided, a new tensor is created.  
+
+    Example:
+
+    .. code-block:: python
+
+        >>> import oneflow as flow
+        >>> flow.normal(mean=0.5, std=flow.arange(1., 6.)).shape
+        oneflow.Size([5])
+
+
+    normal(mean, std, *, generator=None, out=None) -> Tensor
+
+    Returns a tensor of random numbers drawn from separate normal distributions
+    whose mean and standard deviation are given.  
+
+    The :attr:`mean` is a tensor with the mean of  
+    each output element's normal distribution  
+
+    :attr:`std` is a double with the standard deviation of  
+    each output element's normal distribution  
+
+    Args:  
+        mean (Tensor): the tensor of per-element means  
+        std (double): the standard deviation  
+
+    Keyword args:  
+        generator (Generator, optional): Random number generator. Defaults to `oneflow::DefaultGenerator` if not provided.  
+        out (Tensor): The output tensor
+
+    Returns:
+        Tensor: The output tensor, with random normal values.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> import oneflow as flow
+        >>> flow.normal(mean=flow.arange(1., 6.)).shape
+        oneflow.Size([5])
+
+    normal(mean, std, size, *, out=None, placement=None, sbp=None, generator=None, dtype=None, device=None, requires_grad=False) -> Tensor
     Returns a tensor of random numbers drawn from separate normal distributions whose mean and standard deviation are given.
 
     Args: