def merge(ops):
# Simple merge if all X are the same.
if all(o.X is ops[0].X for o in ops):
A, A_sigr = SigMerger.merge([o.A for o in ops])
Y, Y_sigr = SigMerger.merge([o.Y for o in ops])
return (op.DotInc(A, ops[0].X,
Y), Merger.merge_dicts(A_sigr, Y_sigr))
assert all(o1.X is not o2.X for i, o1 in enumerate(ops)
for o2 in ops[i + 1:])
# BSR merge if X differ
X, X_sigr = SigMerger.merge([o.X for o in ops])
Y, Y_sigr = SigMerger.merge([o.Y for o in ops])
# Construct sparse A representation
data = np.array([o.A.initial_value for o in ops], dtype=rc.float_dtype)
if data.ndim == 1:
raise NotImplementedError("A.ndim should be > 2")
elif data.ndim == 2:
raise NotImplementedError("A.ndim should be > 2")
indptr = np.arange(len(ops) + 1, dtype=rc.int_dtype)
indices = np.arange(len(ops), dtype=rc.int_dtype)
name = "bsr_merged<{first}, ..., {last}>".format(first=ops[0].A.name,
last=ops[-1].A.name)
readonly = all([o.A.readonly for o in ops])
A = Signal(data, name=name, readonly=readonly)
A_sigr = {}
for i, s in enumerate([o.A for o in ops]):
A_sigr[s] = Signal(
data[i],
name="%s[%i]" % (s.name, i),
base=A,
offset=i * A.itemsize * np.prod(A.shape[1:]),
)
assert np.allclose(s.initial_value,
A_sigr[s].initial_value,
atol=0,
rtol=0,
equal_nan=True)
assert s.shape == A_sigr[s].shape or (s.shape == () and
A_sigr[s].shape == (1, 1))
reshape = op.reshape_dot(
ops[0].A.initial_value,
ops[0].X.initial_value,
ops[0].Y.initial_value,
tag=ops[0].tag,
)
return (
op.BsrDotInc(A,
X,
Y,
indices=indices,
indptr=indptr,
reshape=reshape),
Merger.merge_dicts(X_sigr, Y_sigr, A_sigr),
)
def merge(ops):
# Simple merge if all X are the same.
if all(o.X is ops[0].X for o in ops):
A, A_sigr = SigMerger.merge([o.A for o in ops])
Y, Y_sigr = SigMerger.merge([o.Y for o in ops])
return (operator.DotInc(A, ops[0].X,
Y), Merger.merge_dicts(A_sigr, Y_sigr))
assert all(o1.X is not o2.X for i, o1 in enumerate(ops)
for o2 in ops[i + 1:])
# BSR merge if X differ
X, X_sigr = SigMerger.merge([o.X for o in ops])
Y, Y_sigr = SigMerger.merge([o.Y for o in ops])
# Construct sparse A representation
data = np.array([o.A.initial_value for o in ops])
if data.ndim == 1:
data = data.reshape((data.size, 1, 1))
elif data.ndim == 2:
data = data.reshape(data.shape + (1, ))
indptr = np.arange(len(ops) + 1, dtype=int)
indices = np.arange(len(ops), dtype=int)
name = 'bsr_merged<{first}, ..., {last}>'.format(first=ops[0].A.name,
last=ops[-1].A.name)
readonly = all([o.A.readonly for o in ops])
A = Signal(data, name=name, readonly=readonly)
A_sigr = {}
for i, s in enumerate([o.A for o in ops]):
A_sigr[s] = Signal(data[i],
name="%s[%i]" % (s.name, i),
base=A,
offset=i * A.itemsize * np.prod(A.shape[1:]))
assert np.all(s.initial_value == A_sigr[s].initial_value)
assert s.shape == A_sigr[s].shape or (s.shape == () and
A_sigr[s].shape == (1, 1))
reshape = operator.reshape_dot(ops[0].A.initial_value,
ops[0].X.initial_value,
ops[0].Y.initial_value,
tag=ops[0].tag)
return (operator.BsrDotInc(A,
X,
Y,
indices=indices,
indptr=indptr,
reshape=reshape),
Merger.merge_dicts(X_sigr, Y_sigr, A_sigr))