def testScatterWeightedSum( self, num_args, first_dim, index_dim, extra_dims, ind_type, gc, dc): ins = ['data', 'w0', 'indices'] for i in range(1, num_args + 1): ins.extend(['x' + str(i), 'w' + str(i)]) op = core.CreateOperator( 'ScatterWeightedSum', ins, ['data'], device_option=gc) def ref(d, w0, ind, *args): r = d.copy() for i in ind: r[i] *= w0 for i in range(0, len(args), 2): x = args[i] w = args[i+1] for i, j in enumerate(ind): r[j] += w * x[i] return [r] d = rand_array(first_dim, *extra_dims) ind = np.random.randint(0, first_dim, index_dim).astype(ind_type) # ScatterWeightedSumOp only supports w0=1.0 in CUDAContext if(gc == hu.gpu_do): w0 = np.array(1.0).astype(np.float32) else: w0 = rand_array() inputs = [d, w0, ind] for _ in range(1, num_args + 1): x = rand_array(index_dim, *extra_dims) w = rand_array() inputs.extend([x,w]) self.assertReferenceChecks(gc, op, inputs, ref, threshold=1e-3)
def testScatterWeightedSum(self, num_args, first_dim, index_dim, extra_dims, ind_type, gc, dc): ins = ['data', 'w0', 'indices'] for i in range(1, num_args + 1): ins.extend(['x' + str(i), 'w' + str(i)]) op = core.CreateOperator('ScatterWeightedSum', ins, ['data'], device_option=gc) def ref(d, w0, ind, *args): r = d.copy() for i in ind: r[i] *= w0 for i in range(0, len(args), 2): x = args[i] w = args[i + 1] for i, j in enumerate(ind): r[j] += w * x[i] return [r] d = rand_array(first_dim, *extra_dims) ind = np.random.randint(0, first_dim, index_dim).astype(ind_type) # ScatterWeightedSumOp only supports w0=1.0 in CUDAContext if (gc == hu.gpu_do): w0 = np.array(1.0).astype(np.float32) else: w0 = rand_array() inputs = [d, w0, ind] for _ in range(1, num_args + 1): x = rand_array(index_dim, *extra_dims) w = rand_array() inputs.extend([x, w]) self.assertReferenceChecks(gc, op, inputs, ref, threshold=1e-3)
def testScatterWeightedSum(self): for num_args in [1, 2]: ins = ['data', 'w0', 'indices'] for i in range(1, num_args + 1): ins.extend(['x' + str(i), 'w' + str(i)]) op = core.CreateOperator('ScatterWeightedSum', ins, ['data']) for first_dim, index_dim, extra_dims in self.test_configs(): d = rand_array(first_dim, *extra_dims) ind = np.random.randint(0, first_dim, index_dim).astype(np.int32) w0 = rand_array() r = d.copy() for i in ind: r[i] *= w0 # forward workspace.FeedBlob('data', d) workspace.FeedBlob('w0', w0) workspace.FeedBlob('indices', ind) for inp in range(1, num_args + 1): w = rand_array() x = rand_array(index_dim, *extra_dims) workspace.FeedBlob('x' + str(inp), x) workspace.FeedBlob('w' + str(inp), w) for i, j in enumerate(ind): r[j] += w * x[i] workspace.RunOperatorOnce(op) out = workspace.FetchBlob('data') np.testing.assert_allclose(out, r, rtol=1e-5)
def testScatterAssign( self, first_dim, index_dim, extra_dims, ind_type, gc, dc): op = core.CreateOperator('ScatterAssign', ['data', 'indices', 'slices'], ['data']) def ref(d, ind, x): r = d.copy() r[ind] = x return [r] # let's have indices unique if first_dim < index_dim: first_dim, index_dim = index_dim, first_dim d = rand_array(first_dim, *extra_dims) ind = np.random.choice(first_dim, index_dim, replace=False).astype(ind_type) x = rand_array(index_dim, *extra_dims) self.assertReferenceChecks(gc, op, [d, ind, x], ref, threshold=1e-3)
def testScatterAssign( self, first_dim, index_dim, extra_dims, data_type, ind_type, gc, dc): op = core.CreateOperator('ScatterAssign', ['data', 'indices', 'slices'], ['data']) def ref(d, ind, x): r = d.copy() r[ind] = x return [r] # let's have indices unique if first_dim < index_dim: first_dim, index_dim = index_dim, first_dim d = (rand_array(first_dim, *extra_dims) * 10).astype(data_type) ind = np.random.choice(first_dim, index_dim, replace=False).astype(ind_type) x = (rand_array(index_dim, *extra_dims) * 10).astype(data_type) self.assertReferenceChecks(gc, op, [d, ind, x], ref, threshold=1e-3)
def testPartition(self): for main_dims, parts, main_type, extra_ins, pack in self.test_configs( ): ins = ['in' + str(i) for i in range(1 + len(extra_ins))] outs = [ 'in{}_p{}'.format(i, j) for i in range(parts) for j in range(1 + len(extra_ins)) ] op = core.CreateOperator('Partition', ins, outs, pack_first_input=(1 if pack else 0)) x = [] for i, (dims, t) in enumerate([((), main_type)] + extra_ins): if t in [np.float32, np.float64]: d = rand_array(*(main_dims + dims)) else: d = np.random.randint(-100, 100, (main_dims + dims)) d = d.astype(t) workspace.FeedBlob(ins[i], d) x.append(d) def sharding(x): # numpy has proper modulo op that yields non-negative results shards = (x[0] % parts).reshape([-1]) out = [] for i in range(parts): for ind, v in enumerate(x): suffix_shape = v.shape[len(x[0].shape):] accum = [] data = v.reshape((-1, ) + suffix_shape) if pack and ind == 0: data = data // parts for j, s in enumerate(shards): if s == i: accum.append(data[j]) def join(a): if not a: return np.empty(shape=(0, ) + suffix_shape) return np.stack(a) out.append(join(accum)) return out workspace.RunOperatorOnce(op) ref = sharding(x) print(x) print(ref) for name, expected in zip(outs, ref): np.testing.assert_array_equal(expected, workspace.FetchBlob(name))
def testPartition(self): for main_dims, parts, main_type, extra_ins, pack in self.test_configs(): ins = ['in' + str(i) for i in range(1 + len(extra_ins))] outs = [ 'in{}_p{}'.format(i, j) for i in range(parts) for j in range(1 + len(extra_ins)) ] op = core.CreateOperator( 'Partition', ins, outs, pack_first_input=(1 if pack else 0)) x = [] for i, (dims, t) in enumerate([((), main_type)] + extra_ins): if t in [np.float32, np.float64]: d = rand_array(*(main_dims + dims)) else: d = np.random.randint(-100, 100, (main_dims + dims)) d = d.astype(t) workspace.FeedBlob(ins[i], d) x.append(d) def sharding(x): # numpy has proper modulo op that yields non-negative results shards = (x[0] % parts).reshape([-1]) out = [] for i in range(parts): for ind, v in enumerate(x): suffix_shape = v.shape[len(x[0].shape):] accum = [] data = v.reshape((-1, ) + suffix_shape) if pack and ind == 0: data = data // parts for j, s in enumerate(shards): if s == i: accum.append(data[j]) def join(a): if not a: return np.empty(shape=(0, ) + suffix_shape) return np.stack(a) out.append(join(accum)) return out workspace.RunOperatorOnce(op) ref = sharding(x) print(x) print(ref) for name, expected in zip(outs, ref): np.testing.assert_array_equal( expected, workspace.FetchBlob(name) )
def testScatterAssign(self): op = core.CreateOperator('ScatterAssign', ['data', 'indices', 'slices'], ['data']) for first_dim, index_dim, extra_dims in self.test_configs(): # let's have indices unique if first_dim < index_dim: first_dim, index_dim = index_dim, first_dim for dtype in [np.int32, np.int64]: d = rand_array(first_dim, *extra_dims) ind = np.random.choice(first_dim, index_dim, replace=False).astype(dtype) x = rand_array(index_dim, *extra_dims) r = d.copy() r[ind] = x # forward workspace.FeedBlob('data', d) workspace.FeedBlob('indices', ind) workspace.FeedBlob('slices', x) workspace.RunOperatorOnce(op) out = workspace.FetchBlob('data') np.testing.assert_allclose(out, r, rtol=1e-3)
def testLengthsPartition(self): for main_dims, parts, main_type, extra_ins, pack in self.test_configs( ): # For LengthsSharding only 1-D tensors supported as a first input if len(main_dims) > 1: continue ins = ['in' + str(i) for i in range(2 + len(extra_ins))] outs = [ 'in{}_p{}'.format(j, i) for i in range(parts) for j in range(2 + len(extra_ins)) ] op = core.CreateOperator('LengthsPartition', ins, outs, pack_first_input=(1 if pack else 0)) x = [] for i, (dims, t) in enumerate([((), main_type)] + extra_ins): if t in [np.float32, np.float64]: d = rand_array(*(main_dims + dims)) else: d = np.random.randint(-100, 100, (main_dims + dims)) d = d.astype(t) workspace.FeedBlob(ins[i + 1], d) x.append(d) # Randomly generate length tensor as well elements = np.random.randint(2, 10) lengths = [] total_length = 0 for _ in range(elements - 1): lengths.append(np.random.randint(main_dims[0] - total_length)) total_length += lengths[-1] lengths.append(main_dims[0] - total_length) workspace.FeedBlob(ins[0], np.array(lengths, dtype=np.int32)) def sharding(x): # numpy has proper modulo op that yields non-negative results shards = (x[0] % parts).reshape([-1]) out = [] for i in range(parts): idx = 0 sharded_lengths = np.zeros(elements) for ind, length in enumerate(lengths): for _ in range(length): if shards[idx] == i: sharded_lengths[ind] += 1 idx += 1 out.append(sharded_lengths) for ind, v in enumerate(x): suffix_shape = v.shape[len(x[0].shape):] accum = [] data = v.reshape((-1, ) + suffix_shape) if pack and ind == 0: data = data // parts for j, s in enumerate(shards): if s == i: accum.append(data[j]) def join(a): if not a: return np.empty(shape=(0, ) + suffix_shape) return np.stack(a) out.append(join(accum)) return out workspace.RunOperatorOnce(op) ref = sharding(x) for name, expected in zip(outs, ref): np.testing.assert_array_equal(expected, workspace.FetchBlob(name))
def testLengthsPartition(self): for main_dims, parts, main_type, extra_ins, pack in self.test_configs(): # For LengthsSharding only 1-D tensors supported as a first input if len(main_dims) > 1: continue ins = ['in' + str(i) for i in range(2 + len(extra_ins))] outs = [ 'in{}_p{}'.format(j, i) for i in range(parts) for j in range(2 + len(extra_ins)) ] op = core.CreateOperator( 'LengthsPartition', ins, outs, pack_first_input=(1 if pack else 0) ) x = [] for i, (dims, t) in enumerate([((), main_type)] + extra_ins): if t in [np.float32, np.float64]: d = rand_array(*(main_dims + dims)) else: d = np.random.randint(-100, 100, (main_dims + dims)) d = d.astype(t) workspace.FeedBlob(ins[i + 1], d) x.append(d) # Randomly generate length tensor as well elements = np.random.randint(2, 10) lengths = [] total_length = 0 for i in range(elements - 1): lengths.append(np.random.randint(main_dims[0] - total_length)) total_length += lengths[-1] lengths.append(main_dims[0] - total_length) workspace.FeedBlob(ins[0], np.array(lengths, dtype=np.int32)) def sharding(x): # numpy has proper modulo op that yields non-negative results shards = (x[0] % parts).reshape([-1]) out = [] for i in range(parts): idx = 0 sharded_lengths = np.zeros(elements) for ind, length in enumerate(lengths): for j in range(length): if shards[idx] == i: sharded_lengths[ind] += 1 idx += 1 out.append(sharded_lengths) for ind, v in enumerate(x): suffix_shape = v.shape[len(x[0].shape):] accum = [] data = v.reshape((-1, ) + suffix_shape) if pack and ind == 0: data = data // parts for j, s in enumerate(shards): if s == i: accum.append(data[j]) def join(a): if not a: return np.empty(shape=(0, ) + suffix_shape) return np.stack(a) out.append(join(accum)) return out workspace.RunOperatorOnce(op) ref = sharding(x) for name, expected in zip(outs, ref): np.testing.assert_array_equal( expected, workspace.FetchBlob(name) )