def test_signal_dict_combine(): minibatch_size = 1 signals = SignalDict(tf.float32, minibatch_size) key = object() assert signals.combine([]) == [] y = signals.combine( [ signals.get_tensor_signal([(0, 3)], key, None, (3, 2), False), signals.get_tensor_signal([(4, 7)], key, None, (3, 2), False), ] ) assert y.key is key assert y._tf_indices is None assert y.shape == (6, 2) assert y.slices == ((0, 3), (4, 7)) y = signals.combine( [ signals.get_tensor_signal([(0, 3)], key, None, (3, 2), False), signals.get_tensor_signal([(3, 6)], key, None, (3, 2), False), ] ) assert y.key is key assert y._tf_indices is None assert y.shape == (6, 2) assert y.slices == ((0, 6),)
def test_signal_dict_scatter(minibatched): minibatch_size = 2 var_size = 19 signals = SignalDict(tf.float32, minibatch_size) key = object() var_key = object() val = np.random.random( (minibatch_size, var_size) if minibatched else (var_size,) ).astype(np.float32) update_shape = (minibatch_size, 4) if minibatched else (4,) pre_slice = np.index_exp[:, :4] if minibatched else np.index_exp[:4] post_slice = np.index_exp[:, 4:] if minibatched else np.index_exp[4:] signals.bases = {key: tf.constant(val), var_key: tf.Variable(val)} x = signals.get_tensor_signal([(0, 4)], key, tf.float32, (4,), minibatched) with pytest.raises(BuildError, match="wrong dtype"): signals.scatter(x, tf.ones(update_shape, dtype=tf.float64)) x_var = signals.get_tensor_signal([(0, 4)], var_key, tf.float32, (4,), minibatched) with pytest.raises(BuildError, match="should not be a Variable"): signals.scatter(x_var, tf.ones(update_shape)) # update signals.scatter(x, tf.ones(update_shape)) y = signals.bases[key] assert np.allclose(y[pre_slice], 1) assert np.allclose(y[post_slice], val[post_slice]) assert signals.write_types["scatter_update"] == 1 # increment, and reshaping val signals.scatter( x, tf.ones((minibatch_size, 2, 2) if minibatched else (2, 2)), mode="inc" ) y = signals.bases[key] assert np.allclose(y[pre_slice], 2) assert np.allclose(y[post_slice], val[post_slice]) assert signals.write_types["scatter_add"] == 1 # recognize assignment to full array x = signals.get_tensor_signal( [(0, var_size)], key, tf.float32, (var_size,), minibatched ) y = tf.ones((minibatch_size, var_size) if minibatched else (var_size,)) signals.scatter(x, y) assert signals.bases[key] is y assert signals.write_types["assign"] == 1
def test_signal_dict_scatter(): minibatch_size = 1 var_size = 19 signals = SignalDict(tf.float32, minibatch_size) sess = tf.InteractiveSession() key = object() val = np.random.randn(var_size, minibatch_size) signals.bases = {key: tf.assign(tf.Variable(val, dtype=tf.float32), val)} x = signals.get_tensor_signal([0, 1, 2, 3], key, tf.float32, (4,), True) with pytest.raises(BuildError): # wrong dtype signals.scatter(x, tf.ones((4,), dtype=tf.float64)) # update signals.scatter(x, tf.ones((4,))) y = sess.run(signals.bases[key]) assert np.allclose(y[:4], 1) assert np.allclose(y[4:], val[4:]) # increment, and reshaping val signals.scatter(x, tf.ones((2, 2)), mode="inc") y = sess.run(signals.bases[key]) assert np.allclose(y[:4], 2) assert np.allclose(y[4:], val[4:]) # recognize assignment to full array x = signals.get_tensor_signal(np.arange(var_size), key, tf.float32, (var_size,), True) y = tf.ones((var_size, 1)) signals.scatter(x, y) assert signals.bases[key].op.type == "Assign" # recognize assignment to strided full array x = signals.get_tensor_signal(np.arange(0, var_size, 2), key, tf.float32, (var_size // 2 + 1,), True) y = tf.ones((var_size // 2 + 1, 1)) signals.scatter(x, y) assert signals.bases[key].op.type == "ScatterUpdate" sess.close()
def test_signal_dict_gather(): minibatch_size = 1 var_size = 19 signals = SignalDict(tf.float32, minibatch_size) sess = tf.InteractiveSession() key = object() val = np.random.randn(var_size, minibatch_size) signals.bases = {key: tf.constant(val, dtype=tf.float32)} x = signals.get_tensor_signal([0, 1, 2, 3], key, tf.float32, (4,), True) # sliced read assert np.allclose(sess.run(signals.gather(x)), val[:4]) # read with reshape x = signals.get_tensor_signal([0, 1, 2, 3], key, tf.float32, (2, 2), True) assert np.allclose(sess.run(signals.gather(x)), val[:4].reshape((2, 2, minibatch_size))) # gather read x = signals.get_tensor_signal([0, 1, 2, 3], key, tf.float32, (4,), True) y = signals.gather(x, force_copy=True) assert "Gather" in y.op.type x = signals.get_tensor_signal([0, 0, 3, 3], key, tf.float32, (4,), True) assert np.allclose(sess.run(signals.gather(x)), val[[0, 0, 3, 3]]) assert "Gather" in y.op.type # reading from full array x = signals.get_tensor_signal(np.arange(var_size), key, tf.float32, (var_size,), True) y = signals.gather(x) assert y is signals.bases[key] # reading from strided full array x = signals.get_tensor_signal(np.arange(0, var_size, 2), key, tf.float32, (var_size // 2 + 1,), True) y = signals.gather(x) assert y.op.type == "StridedSlice" assert y.op.inputs[0] is signals.bases[key] # minibatch dimension x = signals.get_tensor_signal([0, 1, 2, 3], key, tf.float32, (4,), True) assert signals.gather(x).get_shape() == (4, 1) x = signals.get_tensor_signal([0, 1, 2, 3], key, tf.float32, (4,), False) assert signals.gather(x).get_shape() == (4,) sess.close()
def test_signal_dict_gather(minibatched): minibatch_size = 3 var_size = 19 signals = SignalDict(tf.float32, minibatch_size) key = object() val = np.random.random( (minibatch_size, var_size) if minibatched else (var_size,) ).astype(np.float32) gathered_val = val[:, :4] if minibatched else val[:4] signals.bases = {key: tf.constant(val, dtype=tf.float32)} x = signals.get_tensor_signal([(0, 4)], key, tf.float32, (4,), minibatched) # sliced read assert np.allclose(signals.gather(x), gathered_val) assert signals.read_types["strided_slice"] == 1 # read with reshape x = signals.get_tensor_signal([(0, 4)], key, tf.float32, (2, 2), minibatched) y = signals.gather(x) shape = (minibatch_size, 2, 2) if minibatched else (2, 2) assert y.shape == shape assert np.allclose(y, gathered_val.reshape(shape)) assert signals.read_types["strided_slice"] == 2 # gather read x = signals.get_tensor_signal([(0, 4)], key, tf.float32, (4,), minibatched) y = signals.gather(x, force_copy=True) assert signals.read_types["gather"] == 1 x = signals.get_tensor_signal( ((0, 1), (0, 1), (3, 4), (3, 4)), key, tf.float32, (4,), minibatched ) assert np.allclose( signals.gather(x), val[:, [0, 0, 3, 3]] if minibatched else val[[0, 0, 3, 3]] ) assert signals.read_types["gather"] == 2 # reading from full array x = signals.get_tensor_signal( [(0, var_size)], key, tf.float32, (var_size,), minibatched ) y = signals.gather(x) assert y is signals.bases[key] assert signals.read_types["identity"] == 1 # reading from strided full array x = signals.get_tensor_signal( tuple((i * 2, i * 2 + 1) for i in range(var_size // 2 + 1)), key, tf.float32, (var_size // 2 + 1,), minibatched, ) y = signals.gather(x) assert y is not signals.bases[key] assert signals.read_types["gather"] == 3
def test_get_tensor_signal(): signals = SignalDict(tf.float32, 3) # check that tensor_signal is created correctly key = object() tensor_signal = signals.get_tensor_signal( (0,), key, np.float64, (3, 4), True) assert isinstance(tensor_signal, TensorSignal) assert np.array_equal(tensor_signal.indices, (0,)) assert tensor_signal.key == key assert tensor_signal.dtype == np.float64 assert tensor_signal.shape == (3, 4) assert tensor_signal.minibatch_size == 3 assert tensor_signal.constant == signals.constant assert len(signals) == 0 # check adding signal to sig_map sig = Signal(np.zeros(4)) sig.minibatched = True tensor_signal = signals.get_tensor_signal( np.arange(4), key, np.float64, (2, 2), True, signal=sig) assert len(signals) == 1 assert signals[sig] is tensor_signal assert next(iter(signals)) is sig assert next(iter(signals.values())) is tensor_signal # error if sig shape doesn't match indices with pytest.raises(AssertionError): sig = Signal(np.zeros((2, 2))) sig.minibatched = True signals.get_tensor_signal( np.arange(4), key, np.float64, (2, 2), True, signal=sig) # error if sig size doesn't match given shape with pytest.raises(AssertionError): sig = Signal(np.zeros(4)) sig.minibatched = True signals.get_tensor_signal( np.arange(4), key, np.float64, (2, 3), True, signal=sig) # error if minibatched doesn't match with pytest.raises(AssertionError): sig = Signal(np.zeros(4)) sig.minibatched = False signals.get_tensor_signal( np.arange(4), key, np.float64, (2, 2), True, signal=sig)
def test_tf_indices_nd(ndims): signals = SignalDict(tf.float32, 10) shape = (3, 4, 5)[:ndims] x = tf.ones(shape) * tf.reshape( tf.range(0, 3, dtype=tf.float32), (-1,) + (1,) * (ndims - 1) ) assert x.shape == shape sig = signals.get_tensor_signal([(0, 1), (2, 3)], None, np.float32, shape, False) indices = sig.tf_indices_nd result = tf.gather_nd(x, indices) assert result.shape == (2,) + shape[1:] assert np.allclose(result[0], 0) assert np.allclose(result[1], 2)
def test_signal_dict_combine(): minibatch_size = 1 signals = SignalDict(tf.float32, minibatch_size) key = object() assert signals.combine([]) == [] y = signals.combine( [signals.get_tensor_signal([0, 1, 2], key, None, (3, 2), False), signals.get_tensor_signal([4, 5, 6], key, None, (3, 2), False)]) assert y.key is key assert y._tf_indices is None assert y.shape == (6, 2) assert np.all(y.indices == [0, 1, 2, 4, 5, 6])