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 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)
)
