Browse Source

docs(mge/functional): update functional.ones docstring

tags/v1.7.2.m1
Asthestarsfalll 3 years ago
parent
commit
409aa4891f
1 changed files with 26 additions and 19 deletions
  1. +26
    -19
      imperative/python/megengine/functional/tensor.py

+ 26
- 19
imperative/python/megengine/functional/tensor.py View File

@@ -6,7 +6,7 @@
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
from typing import Iterable, Optional, Sequence, Union
from typing import Iterable, Optional, Sequence, Union, Tuple

import numpy as np

@@ -148,33 +148,40 @@ def full(
return broadcast_to(x, shape)


def ones(shape, dtype="float32", device=None) -> Tensor:
r"""Returns a ones tensor with given shape.
def ones(
shape: Union[int, Tuple[int, ...]],
*,
dtype="float32",
device: Optional[CompNode] = None
) -> Tensor:
r"""Returns a new tensor having a specified shape and filled with ones.

Args:
shape: a list, tuple or integer defining the shape of the output tensor.
dtype: the desired data type of the output tensor. Default: ``float32``.
device: the desired device of the output tensor. Default: if ``None``,
use the default device (see :func:`~.megengine.get_default_device`).
shape (int or sequence of ints): the shape of the output tensor.
Keyword args:
dtype (:attr:`.Tensor.dtype`): output tensor data type. Default: ``float32``.
device (:attr:`.Tensor.device`): device on which to place the created tensor. Default: ``None``.

Returns:
output tensor.
a tensor containing ones.

Examples:

.. testcode::
>>> megengine.functional.ones(5)
Tensor([1. 1. 1. 1. 1.], device=xpux:0)

import megengine.functional as F

out = F.ones((2, 1))
print(out.numpy())

Outputs:

.. testoutput::
>>> megengine.functional.ones((5, ), dtype='int32')
Tensor([1 1 1 1 1], dtype=int32, device=xpux:0)
>>> megengine.functional.ones((2, 2))
Tensor([[1. 1.]
[1. 1.]], device=xpux:0)
>>> megengine.functional.ones([2, 1])
Tensor([[1.]
[1.]], device=xpux:0)

[[1.]
[1.]]
"""
return full(shape, 1.0, dtype=dtype, device=device)



Loading…
Cancel
Save