def split(x, sizes_or_sections, dim):
if isinstance(sizes_or_sections, list) and len(sizes_or_sections) == 1:
assert len(x) == sizes_or_sections[0]
return [x]
if isinstance(sizes_or_sections, (np.ndarray, list)):
sizes_or_sections1 = tuple(np.cumsum(sizes_or_sections)[:-1])
return nd.split_v2(x, sizes_or_sections1, axis=dim)
def split(x, sizes_or_sections, dim):
if isinstance(sizes_or_sections, list) and len(sizes_or_sections) == 1:
assert len(x) == sizes_or_sections[0]
return [x]
if MX_VERSION.version[0] == 1 and MX_VERSION.version[1] >= 5:
if isinstance(sizes_or_sections, (np.ndarray, list)):
sizes_or_sections1 = tuple(np.cumsum(sizes_or_sections)[:-1])
return nd.split_v2(x, sizes_or_sections1, axis=dim)
if isinstance(sizes_or_sections, list) or isinstance(sizes_or_sections, np.ndarray):
# Old MXNet doesn't support split with different section sizes.
np_arr = x.asnumpy()
indices = np.cumsum(sizes_or_sections)[:-1]
res = np.split(np_arr, indices, axis=dim)
return [tensor(arr, dtype=x.dtype) for arr in res]
else:
return nd.split(x, sizes_or_sections, axis=dim)
