def test_op_constant(dtype, diff, sess):
ops = (SimNeurons(LIF(tau_rc=1), Signal(np.zeros(10)), None),
SimNeurons(LIF(tau_rc=2 if diff else 1), Signal(np.zeros(10)),
None))
signals = SignalDict(tf.float32, 1)
const = signals.op_constant(
[op.neurons for op in ops], [op.J.shape[0] for op in ops],
"tau_rc", dtype)
const1 = signals.op_constant(
[op.neurons for op in ops], [op.J.shape[0] for op in ops],
"tau_rc", dtype, ndims=1)
const3 = signals.op_constant(
[op.neurons for op in ops], [op.J.shape[0] for op in ops],
"tau_rc", dtype, ndims=3)
assert const.dtype.base_dtype == dtype
sess.run(tf.variables_initializer(tf.get_collection("constants")),
feed_dict=signals.constant_phs)
x, x1, x3 = sess.run([const, const1, const3])
if diff:
assert np.array_equal(x, [[1]] * 10 + [[2]] * 10)
assert np.array_equal(x[:, 0], x1)
assert np.array_equal(x, x3[..., 0])
else:
assert np.array_equal(x, 1.0)
assert np.array_equal(x, x1)
assert np.array_equal(x, x3)
def test_op_constant(dtype, diff):
ops = (
SimNeurons(LIF(tau_rc=1), Signal(np.zeros(10)), None),
SimNeurons(LIF(tau_rc=2 if diff else 1), Signal(np.zeros(10)), None),
)
signals = SignalDict(tf.float32, 1)
const = signals.op_constant(
[op.neurons for op in ops], [op.J.shape[0] for op in ops], "tau_rc", dtype
)
const1 = signals.op_constant(
[op.neurons for op in ops],
[op.J.shape[0] for op in ops],
"tau_rc",
dtype,
shape=(-1,),
)
const3 = signals.op_constant(
[op.neurons for op in ops],
[op.J.shape[0] for op in ops],
"tau_rc",
dtype,
shape=(1, -1, 1),
)
assert const.dtype.base_dtype == dtype
if diff:
assert np.array_equal(const, [[1.0] * 10 + [2.0] * 10])
assert np.array_equal(const[0], const1)
assert np.array_equal(const, const3[..., 0])
else:
assert np.array_equal(const, 1.0)
assert np.array_equal(const, const1)
assert np.array_equal(const, const3)
def test_constant_gpu(sess):
val = np.random.randn(10, 10).astype(np.int32)
with tf.device("/gpu:0"):
signals = SignalDict(tf.float32, 1)
const = signals.constant(val, cutoff=0)
assert const.dtype == tf.int32
assert "GPU" in const.device.upper()
sess.run(tf.variables_initializer(tf.get_collection("constants")),
feed_dict=signals.constant_phs)
c = sess.run(const)
assert np.allclose(val, c)
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])
def test_constant_gpu():
val = np.random.randn(10, 10).astype(np.int32)
graph = tf.Graph()
with graph.as_default(), graph.device("/gpu:0"):
signals = SignalDict(None, tf.float32, 1)
const = signals.constant(val, cutoff=0)
assert const.dtype == tf.int32
assert "GPU" in const.device.upper()
with tf.Session(graph=graph) as sess:
sess.run(tf.variables_initializer(tf.get_collection("constants")),
feed_dict=signals.constant_phs)
c = sess.run(const)
assert np.allclose(val, c)
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_constant(dtype, sess):
val = np.random.randn(10, 10).astype(np.float64)
signals = SignalDict(tf.float32, 1)
const0 = signals.constant(val, dtype=dtype)
const1 = signals.constant(val, dtype=dtype, cutoff=0)
assert const0.op.type == "Const"
assert const1.op.type == "Identity"
assert const0.dtype == (dtype if dtype else tf.as_dtype(val.dtype))
assert const1.dtype == (dtype if dtype else tf.as_dtype(val.dtype))
sess.run(tf.variables_initializer(tf.get_collection("constants")),
feed_dict=signals.constant_phs)
c0, c1 = sess.run([const0, const1])
assert np.allclose(c0, val)
assert np.allclose(c1, val)
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_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_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_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_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