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)