def testDecideChunks(self):
with option_context() as options:
options.tensor.chunk_store_limit = 64
itemsize = 1
shape = (10, 20, 30)
nsplit = decide_chunk_sizes(shape, None, itemsize)
[
self.assertTrue(all(isinstance(i, Integral) for i in ns))
for ns in nsplit
]
self.assertEqual(shape, tuple(sum(ns) for ns in nsplit))
self.assertGreaterEqual(
options.tensor.chunk_store_limit,
itemsize * np.prod([np.max(a) for a in nsplit]))
itemsize = 2
shape = (20, 30, 40)
nsplit = decide_chunk_sizes(shape, {1: 4}, itemsize)
[
self.assertTrue(all(isinstance(i, Integral) for i in ns))
for ns in nsplit
]
self.assertEqual(shape, tuple(sum(ns) for ns in nsplit))
self.assertGreaterEqual(
options.tensor.chunk_store_limit,
itemsize * np.prod([np.max(a) for a in nsplit]))
itemsize = 2
shape = (20, 30, 40)
nsplit = decide_chunk_sizes(shape, [2, 3], itemsize)
[
self.assertTrue(all(isinstance(i, Integral) for i in ns))
for ns in nsplit
]
self.assertEqual(shape, tuple(sum(ns) for ns in nsplit))
self.assertGreaterEqual(
options.tensor.chunk_store_limit,
itemsize * np.prod([np.max(a) for a in nsplit]))
itemsize = 2
shape = (20, 30, 40)
nsplit = decide_chunk_sizes(shape, [20, 3], itemsize)
[
self.assertTrue(all(isinstance(i, Integral) for i in ns))
for ns in nsplit
]
self.assertEqual(shape, tuple(sum(ns) for ns in nsplit))
self.assertEqual(120,
itemsize *
np.prod([np.max(a) for a in nsplit
])) # 20 * 3 * 1 * 2 exceeds limitation
def _tile_as_shuffle(op):
in_tensor = op.input
tensor = op.outputs[0]
new_shape = op.newshape
shuffle_inputs, shuffle_outputs = [], []
axis_offsets = [[0] + np.cumsum(ns)[:-1].tolist() for ns in in_tensor.nsplits]
max_chunk_size = max(max(tp) for tp in in_tensor.nsplits)
out_nsplits = decide_chunk_sizes(new_shape, max_chunk_size, tensor.dtype.itemsize)
chunk_size_idxes = (range(len(size)) for size in out_nsplits)
for inp in in_tensor.chunks:
offset = tuple(axis_offsets[axis][idx] for axis, idx in enumerate(inp.index))
chunk_op = TensorReshapeMap(_input=inp, _axis_offsets=offset, _oldshape=in_tensor.shape,
_newshape=new_shape, _new_chunk_size=(max_chunk_size,) * len(new_shape),
_dtype=inp.dtype)
shuffle_inputs.append(chunk_op.new_chunk([inp], shape=(np.nan,), index=inp.index))
proxy_chunk = TensorShuffleProxy(dtype=in_tensor.dtype, _tensor_keys=[in_tensor.op.key]) \
.new_chunk(shuffle_inputs, shape=())
for chunk_shape, chunk_idx in izip(itertools.product(*out_nsplits),
itertools.product(*chunk_size_idxes)):
shuffle_key = ','.join(str(o) for o in chunk_idx)
chunk_op = TensorReshapeReduce(_dtype=tensor.dtype, _shuffle_key=shuffle_key)
shuffle_outputs.append(chunk_op.new_chunk([proxy_chunk], shape=chunk_shape, index=chunk_idx))
new_op = op.copy()
return new_op.new_tensors(op.inputs, new_shape, chunks=shuffle_outputs,
nsplits=out_nsplits)