MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / stack

Function stack

imperative/python/megengine/functional/tensor.py:597–629  ·  view source on GitHub ↗

r"""Concats a sequence of tensors along a new axis. The input tensors must have the same shape. Args: inps: input tensors. axis: which axis will be concatenated. device: the device output will be. Default: None Returns: output concatenated tensor. E

(inps, axis=0, device=None)

Source from the content-addressed store, hash-verified

595
596
597def stack(inps, axis=0, device=None):
598 r"""Concats a sequence of tensors along a new axis.
599 The input tensors must have the same shape.
600
601 Args:
602 inps: input tensors.
603 axis: which axis will be concatenated.
604 device: the device output will be. Default: None
605
606 Returns:
607 output concatenated tensor.
608
609 Examples:
610 >>> import numpy as np
611 >>> x1 = Tensor(np.arange(0, 3, dtype=np.float32).reshape((3)))
612 >>> x2 = Tensor(np.arange(6, 9, dtype=np.float32).reshape((3)))
613 >>> out = F.stack([x1, x2], axis=0)
614 >>> out.numpy()
615 array([[0., 1., 2.],
616 [6., 7., 8.]], dtype=float32)
617 """
618 if len(inps) == 1:
619 ret = expand_dims(inps[0], axis=axis)
620 if device is None:
621 return ret
622 else:
623 return copy(ret, device)
624
625 if device is None:
626 device = get_device(inps)
627 device = as_device(device)
628 (result,) = apply(builtin.Stack(axis=axis, comp_node=device.to_c()), *inps)
629 return result
630
631
632def split(inp, nsplits_or_sections, axis=0):

Callers

nothing calls this directly

Calls 6

get_deviceFunction · 0.85
as_deviceFunction · 0.85
to_cMethod · 0.80
expand_dimsFunction · 0.70
copyFunction · 0.70
applyFunction · 0.50

Tested by

no test coverage detected