def test_samples_df(samples_per_chain, chains):
"""Test that the samples df is as expected."""
class Model(Normal):
dv = "y"
features = dict(
x1=dict(transformer=1, prior=dist.Normal(0, 1)),
x2=dict(transformer=2, prior=dist.Normal(0, 1)),
)
config = {
"samples": {
"x1": onp.random.normal(size=(chains, samples_per_chain)),
"x2": onp.random.normal(size=(chains, samples_per_chain)),
"_sigma": onp.ones((chains, samples_per_chain)),
}
}
model = Model.from_dict(config)
df = model.samples_df.reset_index()
assert df["chain"].max() == (chains - 1)
assert df["sample"].max() == (samples_per_chain - 1)
assert df.shape[0] == (samples_per_chain * chains)
assert df.shape[1] == 4 # 2 features, plus two indecies
for x in model.features.keys():
assert np.array_equal(df[x].values, config["samples"][x].flatten())
config_first_chain = config["samples"][x][0, :]
df_first_chain = df.sort_values("sample").loc[lambda x: x.chain == 0,
x].values
assert np.array_equal(df_first_chain, config_first_chain)
def test_ocr_rnn():
length = 5
carry_size = 3
class_count = 4
inputs = np.zeros((1, length, 4))
def rnn():
return Rnn(*GRUCell(carry_size, zeros))
net = Sequential(
rnn(),
rnn(),
rnn(),
lambda x: np.reshape(x, (-1, carry_size)
), # -> same weights for all time steps
Dense(class_count, zeros, zeros),
softmax,
lambda x: np.reshape(x, (-1, length, class_count)))
params = net.init_parameters(PRNGKey(0), inputs)
assert len(params) == 4
cell = params.rnn0.gru_cell
assert len(cell) == 3
assert np.array_equal(np.zeros((7, 3)), cell.update_kernel)
assert np.array_equal(np.zeros((7, 3)), cell.reset_kernel)
assert np.array_equal(np.zeros((7, 3)), cell.compute_kernel)
out = net.apply(params, inputs)
assert np.array_equal(.25 * np.ones((1, 5, 4)), out)
def test_Batched():
out_dim = 1
@parametrized
def unbatched_dense(input):
kernel = parameter((out_dim, input.shape[-1]), ones)
bias = parameter((out_dim, ), ones)
return np.dot(kernel, input) + bias
batch_size = 4
unbatched_params = unbatched_dense.init_parameters(np.zeros(2),
key=PRNGKey(0))
out = unbatched_dense.apply(unbatched_params, np.ones(2))
assert np.array([3.]) == out
dense_apply = vmap(unbatched_dense.apply, (None, 0))
out_batched_ = dense_apply(unbatched_params, np.ones((batch_size, 2)))
assert np.array_equal(np.stack([out] * batch_size), out_batched_)
dense = Batched(unbatched_dense)
params = dense.init_parameters(np.ones((batch_size, 2)), key=PRNGKey(0))
assert_parameters_equal((unbatched_params, ), params)
out_batched = dense.apply(params, np.ones((batch_size, 2)))
assert np.array_equal(out_batched_, out_batched)
def test_reduce_scalar_jit():
from mpi4jax import reduce
arr = rank
res = jax.jit(lambda x: reduce(x, op=MPI.SUM, root=0)[0])(arr)
if rank == 0:
assert jnp.array_equal(res, sum(range(size)))
else:
assert jnp.array_equal(res, arr)
def test_reduce_scalar():
from mpi4jax import reduce
arr = rank
res, _ = reduce(arr, op=MPI.SUM, root=0)
if rank == 0:
assert jnp.array_equal(res, sum(range(size)))
else:
assert jnp.array_equal(res, arr)
def test_allreduce_scalar_jit():
from mpi4jax import allreduce
arr = 1
_arr = 1
res = jax.jit(lambda x: allreduce(x, op=MPI.SUM)[0])(arr)
assert jnp.array_equal(res, arr * size)
assert jnp.array_equal(_arr, arr)
def test_filename_load_save_var_collection(self):
a = objax.nn.Conv2D(16, 16, 3)
b = objax.nn.Conv2D(16, 16, 3)
self.assertFalse(jn.array_equal(a.w.value, b.w.value))
with tempfile.NamedTemporaryFile('wb') as f:
objax.io.save_var_collection(f.name, a.vars())
objax.io.load_var_collection(f.name, b.vars())
self.assertEqual(a.w.value.dtype, b.w.value.dtype)
self.assertTrue(jn.array_equal(a.w.value, a.w.value))
def test_allreduce_jit():
from mpi4jax import allreduce
arr = jnp.ones((3, 2))
_arr = arr.copy()
res = jax.jit(lambda x: allreduce(x, op=MPI.SUM)[0])(arr)
assert jnp.array_equal(res, arr * size)
assert jnp.array_equal(_arr, arr)
def test_allreduce():
from mpi4jax import allreduce
arr = jnp.ones((3, 2))
_arr = arr.copy()
res, token = allreduce(arr, op=MPI.SUM)
assert jnp.array_equal(res, arr * size)
assert jnp.array_equal(_arr, arr)
def test_random():
"""Test random singleton"""
random.set_key()
random.set_key(0)
assert jnp.array_equal(jnp.array([0, 0]), random.get_key())
expected = jnp.array([2718843009, 1272950319], dtype=jnp.uint32)
assert jnp.array_equal(random.generate_key(), expected)
expected = jnp.array([4146024105, 967050713], dtype=jnp.uint32)
assert jnp.array_equal(random.get_key(), expected)
def test_gather_scalar_jit():
from mpi4jax import gather
arr = rank
res = jax.jit(lambda x: gather(x, root=0)[0])(arr)
if rank == 0:
assert jnp.array_equal(res, jnp.arange(size))
else:
assert jnp.array_equal(res, arr)
def test_gather_scalar():
from mpi4jax import gather
arr = rank
res, _ = gather(arr, root=0)
if rank == 0:
assert jnp.array_equal(res, jnp.arange(size))
else:
assert jnp.array_equal(res, arr)
def test_allreduce_scalar():
from mpi4jax import allreduce
arr = 1
_arr = 1
res, token = allreduce(arr, op=MPI.SUM)
assert jnp.array_equal(res, arr * size)
assert jnp.array_equal(_arr, arr)
def test_scan(self):
def cell(carry, x):
return jn.array([2]) * carry * x, jn.array([3]) * carry * x
carry = jn.array([8., 8.])
output = jn.array([[3., 3.], [6., 6.], [12., 12.]])
test_carry, test_output = objax.functional.scan(
cell, jn.ones((2, )), jn.ones((3, )))
self.assertTrue(jn.array_equal(carry, test_carry))
self.assertTrue(jn.array_equal(output, test_output))
def test_file_load_var_collection_rename(self):
a = objax.nn.Conv2D(16, 16, 3)
b = objax.nn.Conv2D(16, 16, 3)
self.assertFalse(jn.array_equal(a.w.value, b.w.value))
with io.BytesIO() as f:
objax.io.save_var_collection(f, a.vars().rename(objax.util.Renamer({'(Conv2D)': '(MyConv2D)'})))
f.seek(0)
objax.io.load_var_collection(f, b.vars(), renamer=objax.util.Renamer({'(MyConv2D)': '(Conv2D)'}))
self.assertEqual(a.w.value.dtype, b.w.value.dtype)
self.assertTrue(jn.array_equal(a.w.value, a.w.value))
def test_file_load_save_var_collection(self):
a = objax.nn.Conv2D(16, 16, 3)
b = objax.nn.Conv2D(16, 16, 3)
self.assertFalse(jn.array_equal(a.w.value, b.w.value))
with io.BytesIO() as f:
objax.io.save_var_collection(f, a.vars())
f.seek(0)
objax.io.load_var_collection(f, b.vars())
self.assertEqual(a.w.value.dtype, b.w.value.dtype)
self.assertTrue(jn.array_equal(a.w.value, a.w.value))
def test_reduce():
from mpi4jax import reduce
arr = jnp.ones((3, 2)) * rank
res, _ = reduce(arr, op=MPI.SUM, root=0)
if rank == 0:
assert jnp.array_equal(res, jnp.ones((3, 2)) * sum(range(size)))
else:
assert jnp.array_equal(res, arr)
def test_parallel_concat_multi_output(self):
"""Parallel inference (concat reduction) for multiple outputs."""
f = objax.nn.Linear(3, 4)
x = objax.random.normal((96, 3))
y = f(x)
fp = objax.Parallel(lambda x: [f(x), f(-x)], vc=f.vars())
with fp.vars().replicate():
z1, z2 = fp(x)
self.assertTrue(jn.array_equal(z1, y))
self.assertTrue(jn.array_equal(z2, -y))
def test_gather():
from mpi4jax import gather
arr = jnp.ones((3, 2)) * rank
res, _ = gather(arr, root=0)
if rank == 0:
for p in range(size):
assert jnp.array_equal(res[p], jnp.ones((3, 2)) * p)
else:
assert jnp.array_equal(res, arr)
def test_allreduce_vmap():
from mpi4jax import allreduce
arr = jnp.ones((3, 2))
_arr = arr.copy()
res = jax.vmap(lambda x: allreduce(x, op=MPI.SUM)[0],
in_axes=0,
out_axes=0)(arr)
assert jnp.array_equal(res, arr * size)
assert jnp.array_equal(_arr, arr)
def test_allreduce_jit_deprecated():
from mpi4jax import Allreduce
arr = jnp.ones((3, 2))
_arr = arr.copy()
with pytest.warns(UserWarning, match="deprecated"):
res = jax.jit(lambda x: Allreduce(x, op=MPI.SUM)[0])(arr)
assert jnp.array_equal(res, arr * size)
assert jnp.array_equal(_arr, arr)
def test_gather_jit():
from mpi4jax import gather
arr = jnp.ones((3, 2)) * rank
res = jax.jit(lambda x: gather(x, root=0)[0])(arr)
if rank == 0:
for p in range(size):
assert jnp.array_equal(res[p], jnp.ones((3, 2)) * p)
else:
assert jnp.array_equal(res, arr)
def test_allreduce():
arr = np.ones((3, 2))
_arr = arr.copy()
res = Allreduce(arr, op=MPI.SUM)
assert np.array_equal(res, arr * size)
assert np.array_equal(_arr, arr)
res = jax.jit(lambda x: Allreduce(x, op=MPI.SUM))(arr)
assert np.array_equal(res, arr * size)
assert np.array_equal(_arr, arr)
def test_allreduce_jvp():
from mpi4jax import allreduce
arr = jnp.ones((3, 2))
_arr = arr.copy()
res, jvp = jax.jvp(lambda x: allreduce(x, op=MPI.SUM)[0], (arr,), (_arr,))
expected, _ = allreduce(arr, op=MPI.SUM)
assert jnp.array_equal(expected, res)
expected, _ = allreduce(_arr, op=MPI.SUM)
assert jnp.array_equal(expected, jvp)
def test_allreduce_vjp():
from mpi4jax import allreduce
arr = jnp.ones((3, 2))
_arr = arr.copy()
res, vjp_fun = jax.vjp(lambda x: allreduce(x, op=MPI.SUM)[0], arr)
(vjp, ) = vjp_fun(_arr)
expected, _ = allreduce(arr, op=MPI.SUM)
assert jnp.array_equal(expected, res)
assert jnp.array_equal(_arr, vjp)
def test_sendrecv_scalar():
from mpi4jax import Sendrecv
arr = 1 * rank
_arr = arr
other = 1 - rank
res, token = Sendrecv(arr, arr, source=other, dest=other)
assert jnp.array_equal(res, jnp.ones_like(arr) * other)
assert jnp.array_equal(_arr, arr)
def test_sendrecv():
from mpi4jax import Sendrecv
arr = jnp.ones((3, 2)) * rank
_arr = arr.copy()
other = 1 - rank
res, token = Sendrecv(arr, arr, source=other, dest=other)
assert jnp.array_equal(res, jnp.ones_like(arr) * other)
assert jnp.array_equal(_arr, arr)
def test_sendrecv_scalar_jit():
from mpi4jax import sendrecv
arr = 1 * rank
_arr = arr
other = 1 - rank
res = jax.jit(lambda x, y: sendrecv(x, y, source=other, dest=other)[0])(arr, arr)
assert jnp.array_equal(res, jnp.ones_like(arr) * other)
assert jnp.array_equal(_arr, arr)
def test_L2Regularized_sequential():
loss = Sequential(Dense(1, ones, ones), relu, Dense(1, ones, ones), sum)
reg_loss = L2Regularized(loss, scale=2)
inputs = np.ones(1)
params = reg_loss.init_parameters(PRNGKey(0), inputs)
assert np.array_equal(np.ones((1, 1)), params.model.dense0.kernel)
assert np.array_equal(np.ones((1, 1)), params.model.dense1.kernel)
reg_loss_out = reg_loss.apply(params, inputs)
assert 7 == reg_loss_out
def test_external_param_sharing():
layer = Dense(2, zeros, zeros)
shared_net = Sequential(layer, layer)
inputs = np.zeros((1, 2))
params = shared_net.init_parameters(inputs, key=PRNGKey(0))
assert_parameters_equal(((np.zeros((2, 2)), np.zeros(2)), ), params)
out = shared_net.apply(params, inputs)
assert np.array_equal(np.zeros((1, 2)), out)
out = shared_net.apply(params, inputs, jit=True)
assert np.array_equal(np.zeros((1, 2)), out)
