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

Function split

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

r"""Splits the input tensor into several smaller tensors. When nsplits_or_sections is int, the last tensor may be smaller than others. Args: inp: input tensor. nsplits_or_sections: number of sub tensors or sections information list. axis: which axis will be splited.

(inp, nsplits_or_sections, axis=0)

Source from the content-addressed store, hash-verified

630
631
632def split(inp, nsplits_or_sections, axis=0):
633 r"""Splits the input tensor into several smaller tensors.
634 When nsplits_or_sections is int, the last tensor may be smaller than others.
635
636 Args:
637 inp: input tensor.
638 nsplits_or_sections: number of sub tensors or sections information list.
639 axis: which axis will be splited.
640
641 Returns:
642 output tensor list.
643
644 Examples:
645 >>> import os
646 >>> import numpy as np
647 >>> x = Tensor(np.random.random((10, 20)), dtype=np.float32)
648 >>> y = F.split(x, 3)
649 >>> z = F.split(x, [6, 17], axis=1)
650 >>> print([i.numpy().shape for i in y])
651 [(4, 20), (3, 20), (3, 20)]
652 >>> print([i.numpy().shape for i in z])
653 [(10, 6), (10, 11), (10, 3)]
654 """
655
656 return split_cpp(inp, nsplits_or_sections, axis)
657
658
659def _get_idx(index, axis):

Callers 1

rollFunction · 0.70

Calls

no outgoing calls

Tested by

no test coverage detected