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, # Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an # software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # "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 import numpy as np


@@ -148,33 +148,40 @@ def full(
return broadcast_to(x, shape) 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: 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: Returns:
output tensor.
a tensor containing ones.


Examples: 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) return full(shape, 1.0, dtype=dtype, device=device)




Loading…
Cancel
Save