From 824b1f71043b58bfae53952ec4ba805b09d2772c Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Fri, 7 Jul 2023 14:49:19 +0800 Subject: [PATCH] docs(functional/complex): add basic docstring for complex apis GitOrigin-RevId: dc93eacdaf8461edf279ec026b4e73c169a8d3f1 --- .../python/megengine/functional/tensor.py | 69 +++++++++++++++++++ imperative/src/impl/basic_operators.cpp | 3 +- .../imperative/transformations/complex.h | 12 +++- 3 files changed, 81 insertions(+), 3 deletions(-) diff --git a/imperative/python/megengine/functional/tensor.py b/imperative/python/megengine/functional/tensor.py index 3bb6b96b9..cb9133f81 100755 --- a/imperative/python/megengine/functional/tensor.py +++ b/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) diff --git a/imperative/src/impl/basic_operators.cpp b/imperative/src/impl/basic_operators.cpp index e3c4c034c..77fe06c8c 100644 --- a/imperative/src/impl/basic_operators.cpp +++ b/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 inputs) const -> Args { diff --git a/imperative/src/include/megbrain/imperative/transformations/complex.h b/imperative/src/include/megbrain/imperative/transformations/complex.h index ea328551a..34bab0328 100644 --- a/imperative/src/include/megbrain/imperative/transformations/complex.h +++ b/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); -- GitLab