def test_flatten_blocks(self):
from ipu_sparse_ops import sparse
values = np.array([1, 2, 3, 4, 5, 6, 7, 8])
print(f"shape:{values.shape}")
spec1 = sparse.matmul_spec_from_max(4, [4, 4],
max_non_zeros=values.size,
block_size=1,
dtype=tf.float32)
bs = 2
spec2 = sparse.matmul_spec_from_max(4, [4, 4],
max_non_zeros=values.size // bs,
block_size=bs,
dtype=tf.float32)
# Trying to flatten already flat values returns same array:
assert_equal(values, sparse.flatten_blocks(spec1, values))
# Block-size 1 has no effect:
not_blocks = sparse.unflatten_blocks(spec1, values)
assert_equal(not_blocks, values)
# Test round trip with block size 2:
blocks = sparse.unflatten_blocks(spec2, values)
assert blocks.shape == (bs, bs, bs)
values_rt = sparse.flatten_blocks(spec2, blocks)
assert_equal(values, values_rt)
def test_block_conversions(self):
from ipu_sparse_ops import sparse
a = np.kron([[1, 0], [1, 0]], [[1, 2], [3, 4]])
b = np.kron([[0, 0], [1, 0]], [[4, 4], [4, 4]])
dense = a + b
bs = 2
spec = sparse.matmul_spec_from_max(2 * bs, [1, 2 * bs],
2,
block_size=bs,
dtype=tf.float32)
blocks = np.reshape([1, 2, 3, 4, 5, 6, 7, 8], [2, bs, bs])
t = ([0, 1], [0, 0], blocks)
n = sparse.dense_from_triplets(spec, *t)
assert_equal(dense, n)
# Check that mask from dense and mask from triplets
# return the same result:
mask_dense = np.zeros_like(dense)
mask_dense[np.nonzero(dense)] = 1
mask_trips = sparse.mask_from_triplets(spec, *t)
assert_equal(mask_dense, mask_trips)
# Check triplets from dense returns same triplets:
td = sparse.triplets_from_dense(dense, bs)
assert_equal(t[0], td.row_indices)
assert_equal(t[1], td.col_indices)
assert_equal(t[2], td.values)
def test_random_indices(self):
from ipu_sparse_ops import sparse
spec = sparse.matmul_spec_from_max(10, [1, 20], 10, tf.float32)
r, c = sparse.random_indices(spec, None)
print(f"r,c:\n{r}\n{c}")
assert len(r) == 10
assert len(r) == len(c)
assert np.max(r) < spec.input_size
assert np.max(c) < spec.output_size
assert np.min(r) >= 0
assert np.min(c) >= 0
def test_conversions(self):
from ipu_sparse_ops import sparse
m = np.array([[10, 0], [0, 20]])
t = sparse.triplets_from_dense(m)
assert_equal(t[0], [0, 1])
assert_equal(t[1], [0, 1])
assert_equal(t[2], [10, 20])
spec = sparse.matmul_spec_from_max(2, [1, 2], 2, tf.float32)
n = sparse.dense_from_triplets(spec, *t)
assert_equal(n, m)
o = sparse.mask_from_triplets(spec, *t)
assert_equal(o, np.array([[1, 0], [0, 1]]))
def test_representation_round_trip_elements(self):
from ipu_sparse_ops import sparse
bs = 16
block_mask = np.array([[1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1]])
mask = np.kron(block_mask, np.ones(shape=[bs, bs])).astype(int)
n_els = np.count_nonzero(mask)
dense = np.zeros_like(mask)
dense[np.nonzero(mask)] = np.arange(n_els)
opts = {"metaInfoBucketOversizeProportion": 1}
t = sparse.triplets_from_dense(dense)
spec = sparse.matmul_spec_from_max(dense.shape[1], [2, dense.shape[0]],
max_non_zeros=n_els,
block_size=1,
dtype=tf.float32)
r = sparse.representation_from_triplets(spec, *t, opts)
t_rt = sparse.triplets_from_representation(spec, r, opts)
dense_rt = sparse.dense_from_triplets(spec, *t_rt)
assert_equal(dense, dense_rt)
def test_disjoint_random_indices(self):
from ipu_sparse_ops import sparse
spec = sparse.matmul_spec_from_max(10, [1, 20], 10, tf.float32)
ta, tb = sparse.disjoint_random_indices(spec,
size_a=10,
size_b=5,
indices_initialiser_gen=None)
print(f"r,c:\n{ta}\n{tb}")
assert len(ta[0]) == 10
assert len(tb[0]) == 5
# Check the result is disjoint by checking uniqueness of ravelled indices:
shape = (spec.input_size, spec.output_size)
a_flat = np.ravel_multi_index((ta[0], ta[1]), shape)
b_flat = np.ravel_multi_index((tb[0], tb[1]), shape)
all_flat = np.concatenate((a_flat, b_flat), axis=0)
unique = np.unique(all_flat)
assert len(unique) == len(all_flat)
# Check we get an error if total number of indices is too large to be unique:
with pytest.raises(ValueError):
ta, tb = sparse.disjoint_random_indices(
spec, size_a=150, size_b=100, indices_initialiser_gen=None)
def test_representation_round_trip_blocks(self):
from ipu_sparse_ops import sparse
for bs in [4, 8, 16]:
# Create a mask that describes the non-zero block structure:
block_mask = np.array([[1, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 1]])
n_blocks = np.count_nonzero(block_mask)
# From that produce an element-wise mask using a Kronecker product:
mask = np.kron(block_mask, np.ones(shape=[bs, bs])).astype(int)
n_els = np.count_nonzero(mask)
# Make a dense matrix from the element-wise mask and fill with random values:
dense = np.zeros_like(mask, dtype=np.float32)
values = np.random.rand(n_els)
dense[np.nonzero(mask)] = values
# Make the spec for the sparse matmul:
opts = {"metaInfoBucketOversizeProportion": 1}
spec = sparse.matmul_spec_from_max(dense.shape[1],
[2, dense.shape[0]],
max_non_zeros=n_blocks,
block_size=bs,
dtype=tf.float32)
# Make triplets indices from the block mask:
t = sparse.triplets_from_dense(block_mask)
# Then fill in triplet's values by extracting the blocks
# from the dense matrix (this can't be done by reshaping):
t_block = sparse.Triplets(
t.row_indices, t.col_indices,
sparse.blocks_at_indices(t.row_indices, t.col_indices, bs,
dense))
# Convert to on device representation and back and check the
# result is the dense matrix we sytarted with:
r = sparse.representation_from_triplets(spec, *t_block, opts)
t_rt = sparse.triplets_from_representation(spec, r, opts)
dense_rt = sparse.dense_from_triplets(spec, *t_rt)
assert_equal(dense, dense_rt)
# Check triplets from dense returns original triplets:
td = sparse.triplets_from_dense(dense_rt, bs)
assert_equal(t_block.row_indices, td.row_indices)
assert_equal(t_block.col_indices, td.col_indices)
assert_equal(t_block.values, td.values)
def test_device_version_equality_ipu2(self):
from ipu_sparse_ops import sparse
bs = 16
block_mask = np.array([[1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1]])
mask = np.kron(block_mask, np.ones(shape=[bs, bs])).astype(int)
n_els = np.count_nonzero(mask)
dense = np.zeros_like(mask)
dense[np.nonzero(mask)] = np.arange(n_els)
opts = {"metaInfoBucketOversizeProportion": 1}
t = sparse.triplets_from_dense(dense)
spec = sparse.matmul_spec_from_max(dense.shape[1], [2, dense.shape[0]],
max_non_zeros=n_els,
block_size=1,
dtype=tf.float32)
# from device
device_r = sparse.representation_from_triplets(spec,
*t,
opts,
ipu_version=0)
device_t_rt = sparse.triplets_from_representation(spec,
device_r,
opts,
ipu_version=0)
# from version
version_r = sparse.representation_from_triplets(spec,
*t,
opts,
ipu_version=2)
version_t_rt = sparse.triplets_from_representation(spec,
version_r,
opts,
ipu_version=2)
assert_equal(device_r.metainfo_state, version_r.metainfo_state)
assert_equal(device_r.nz_values, version_r.nz_values)
assert_equal(device_t_rt, version_t_rt)
def from_triplets(cls,
hidden_size: int,
input_shape: List[int],
row_indices: List[int],
col_indices: List[int],
values: List[float],
matmul_options: Mapping[str, str],
name: str,
dtype: tf.DType = tf.float32,
use_bias: bool = False,
relu: bool = False,
disable_updating: bool = False,
pooling_type: str = 'NONE'):
"""
Utility factory function to build a 'SparseFcLayer' from a set of triplets (COO format).
E.g. as returned from 'ipu_sparse_ops.sparse.triplets_from_dense'
"""
block_size = sparse.block_size_from_list(values)
spec = sparse.matmul_spec_from_max(hidden_size, input_shape,
len(values), block_size, dtype,
pooling_type)
ns = tf.get_default_graph().get_name_scope()
qualified_name = ns + "/" + name if ns else name
logger.debug(
f"Creating random sparse FC {qualified_name} with spec: {spec}")
triplets = sparse.Triplets(row_indices, col_indices, values)
weights = SparseMatrix(spec,
matmul_options,
triplets,
name=qualified_name)
return cls(weights,
name,
use_bias,
relu,
disable_updating,
pooling_type=pooling_type)