|
|
@@ -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) |
|
|
|
|
|
|
|