docs(functional/complex): add basic docstring for complex apis

GitOrigin-RevId: dc93eacdaf8461edf279ec026b4e73c169a8d3f1

imperative/python/megengine/functional/tensor.py

@@ -425,10 +425,45 @@ def ones_like(inp: Tensor) -> Tensor:
 
 
 def polar(abs: Tensor, angle: Tensor) -> Tensor:
+    r"""Constructs a complex tensor whose elements are Cartesian coordinates
+    corresponding to the polar coordinates with absolute value abs and angle angle.
+
+    Args:
+        abs(Tensor): the absolute value the complex tensor. Must be float.
+        angle(Tensor): the angle of the complex tensor. Must be float.
+
+    Returns:
+        the complex tensor
+
+    Examples:
+        >>> abs = Tensor([1, 2], dtype=np.float32)
+        >>> angle = Tensor([np.pi / 2, 5 * np.pi / 4], dtype=np.float32)
+        >>> z = F.polar(abs, angle)
+        >>> z
+        Tensor([-4.3711e-08+1.j     -1.4142e+00-1.4142j], dtype=complex64, device=xpux:0)
+    """
     return create_complex(abs * cos(angle), abs * sin(angle))
 
 
 def complex(real: Tensor, imag: Tensor) -> Tensor:
+    r"""Constructs a complex tensor with its real part equal to real and its imaginary part equal to imag.
+
+    Args:
+        real(Tensor): the real part of the complex tensor. Must be float.
+        imag(Tensor): the imaginary part of the complex tensor. Must be float.
+
+    Returns:
+        the complex tensor
+
+    Examples:
+        >>> real = Tensor([1, 2], dtype=np.float32)
+        >>> imag = Tensor([3, 4], dtype=np.float32)
+        >>> z = F.complex(real, imag)
+        >>> z
+        Tensor([1.+3.j 2.+4.j], dtype=complex64, device=xpux:0)
+        >>> z.dtype
+        dtype('complex64')
+    """
     if not isinstance(real, Tensor):
         real = Tensor(real)
     if not isinstance(imag, Tensor):
@@ -437,10 +472,44 @@ def complex(real: Tensor, imag: Tensor) -> Tensor:
 
 
 def real(complex: Tensor) -> Tensor:
+    r"""Returns a new tensor containing real values of the complex tensor.
+
+    Args:
+        complex(Tensor) the complex tensor
+
+    Returns:
+        the real part of the complex tensor
+
+    Examples:
+        >>> x=Tensor([0.3100+0.3553j, -0.5445-0.7896j, -1.6492-0.0633j, -0.0638-0.8119j], dtype=np.complex64)
+        
+        >>> F.real(x)
+        Tensor([[ 0.31  ]
+         [-0.5445]
+         [-1.6492]
+         [-0.0638]], device=xpux:0)
+    """
     return get_real(complex)
 
 
 def imag(complex: Tensor) -> Tensor:
+    r"""Returns a new tensor containing imaginary values of the complex tensor.
+
+    Args:
+        complex(Tensor) the complex tensor
+
+    Returns:
+        the imaginary part of the complex tensor
+
+    Examples:
+        >>> x=Tensor([0.3100+0.3553j, -0.5445-0.7896j, -1.6492-0.0633j, -0.0638-0.8119j], dtype=np.complex64)
+        
+        >>> F.imag(x)
+        Tensor([[ 0.3553]
+         [-0.7896]
+         [-0.0633]
+         [-0.8119]], device=xpux:0)
+    """
     return get_imag(complex)
 
 

imperative/src/impl/basic_operators.cpp

@@ -48,7 +48,8 @@ CreateTensor::CreateTensor(Kind kind, CompNode device, TensorLayout layout)
           m_shape(ValueShape::from(layout)),
           m_format(Format::Type::DEFAULT) {
     mgb_assert(
-            layout.is_contiguous() || layout.is_empty(), "layout should be contiguous");
+            layout.is_contiguous() || layout.is_empty(),
+            "layout should be contiguous, got %s", layout.to_string().c_str());
 }
 
 auto CreateTensor::parse(Span<ValueRef> inputs) const -> Args {

imperative/src/include/megbrain/imperative/transformations/complex.h

@@ -13,6 +13,7 @@
 #include "megbrain/imperative/utils/helper.h"
 #include "megbrain/imperative/utils/span.h"
 #include "megbrain/imperative/value.h"
+#include "megbrain/tensor.h"
 #include "megdnn/thin/small_vector.h"
 
 namespace mgb {
@@ -184,17 +185,24 @@ public:
                         f32_host.sub(SubTensorSpec::make_from_layout(real_layout));
                 auto imag_host = f32_host.sub(
                         SubTensorSpec::make_from_offset_elem(imag_layout, 1));
+                // copy into continuous
+                real_layout.init_contiguous_stride();
+                imag_layout.init_contiguous_stride();
+                auto real_value = HostTensorND{create_tensor->device(), real_layout};
+                real_value.copy_from_fixlayout(real_host);
+                auto imag_value = HostTensorND{create_tensor->device(), imag_layout};
+                imag_value.copy_from_fixlayout(imag_host);
                 // create real and imag
                 auto real = imperative::apply(
                         CreateTensor(
                                 create_tensor->kind(), create_tensor->device(),
                                 real_layout),
-                        HostStorage::make(real_host.storage()))[0];
+                        HostStorage::make(real_value.storage()))[0];
                 auto imag = imperative::apply(
                         CreateTensor(
                                 create_tensor->kind(), create_tensor->device(),
                                 imag_layout),
-                        HostStorage::make(imag_host.storage()))[0];
+                        HostStorage::make(imag_value.storage()))[0];
                 return {m_complex_type.make(real, imag)};
             } else {
                 return imperative::apply(op, inputs);