def test_log_prob(self):
t0 = T.zeros([], T.float32)
d = Distribution(
T.float32, [2, 3], continuous=True, reparameterized=True,
event_ndims=1, min_event_ndims=0)
d._log_prob = Mock(return_value=t0)
def do_check(given, group_ndims, args):
if group_ndims is None:
ret = d.log_prob(given)
else:
ret = d.log_prob(given, group_ndims)
self.assertIs(ret, t0)
self.assertEqual(d._log_prob.call_args, ((given,) + args, {}))
d._log_prob.reset_mock()
for sample_shape in ([], [1], [1, 1]):
do_check(T.zeros(sample_shape), None, (0, 1))
for group_ndims in (-1, 0, 1):
do_check(T.zeros(sample_shape),
group_ndims,
(group_ndims, 1 + group_ndims))
for group_ndims in (-2, 2):
with pytest.raises(Exception,
match=re.compile('`min_event_ndims - event_ndims <= group_ndims <= .*'
'sample_ndims - event_ndims` does not hold: ',
re.DOTALL)):
_ = d.log_prob(T.zeros(sample_shape), group_ndims)
d._log_prob.reset_mock()
for group_ndims in (-1, 0, 1, 2):
do_check(T.zeros([5, 2, 3]),
group_ndims,
(group_ndims, 1 + group_ndims))
def test_no_element_sparse_tensor(self):
for coord_first in [True, False]:
# construct the no-element sparse tensor
indices_shape = [2, 0] if coord_first else [0, 2]
x = T.sparse.make_sparse(
T.zeros(indices_shape, dtype=T.index_dtype),
T.zeros([0], dtype=T.float32),
coord_first=coord_first,
shape=[3, 4],
)
assert_allclose(x, np.zeros([3, 4]))
self.assertEqual(T.sparse.value_count(x), 0)
def check_shuffling_flow(ctx, spatial_ndims: int, cls):
num_features = 5
for batch_shape in ([2], [2, 3]):
shape = make_conv_shape(batch_shape, num_features, [6, 7,
8][:spatial_ndims])
# test constructor
flow = cls(num_features)
ctx.assertIn(f'num_features={num_features}', repr(flow))
permutation = tk.layers.get_parameter(flow, 'permutation')
inv_permutation = tk.layers.get_parameter(flow, 'inv_permutation')
assert_equal(T.argsort(permutation), inv_permutation)
assert_equal(T.argsort(inv_permutation), permutation)
flow = tk.layers.jit_compile(flow)
# prepare for the answer
x = T.random.randn(shape)
channel_axis = get_channel_axis(spatial_ndims)
expected_y = T.index_select(x, permutation, axis=channel_axis)
assert_equal(
T.index_select(expected_y, inv_permutation, axis=channel_axis),
x,
)
expected_log_det = T.zeros(batch_shape)
# check the flow
flow_standard_check(ctx, flow, x, expected_y, expected_log_det,
T.random.randn(batch_shape))
def test_ReshapeFlow(self):
flow = ReshapeFlow([4, -1], [-1])
self.assertEqual(flow.x_event_shape, [4, -1])
self.assertEqual(flow.y_event_shape, [-1])
self.assertEqual(flow.get_x_event_ndims(), 2)
self.assertEqual(flow.get_y_event_ndims(), 1)
self.assertIn('x_event_shape=[4, -1]', repr(flow))
self.assertIn('y_event_shape=[-1]', repr(flow))
flow = tk.layers.jit_compile(flow)
x = T.random.randn([2, 3, 4, 5])
expected_y = T.reshape_tail(x, 2, [-1])
expected_log_det = T.zeros([2, 3])
flow_standard_check(self, flow, x, expected_y, expected_log_det,
T.random.randn([2, 3]))
with pytest.raises(ValueError,
match='Too many `-1` specified in `x_event_shape`'):
_ = ReshapeFlow([-1, -1], [-1])
with pytest.raises(ValueError,
match='All elements of `x_event_shape` must be '
'positive integers or `-1`'):
_ = ReshapeFlow([-1, -2], [-1])
with pytest.raises(ValueError,
match='Too many `-1` specified in `y_event_shape`'):
_ = ReshapeFlow([-1], [-1, -1])
with pytest.raises(ValueError,
match='All elements of `y_event_shape` must be '
'positive integers or `-1`'):
_ = ReshapeFlow([-1], [-1, -2])
def test_construct(self):
# no observation
net = BayesianNet()
self.assertEqual(len(net), 0)
self.assertEqual(list(net), [])
self.assertEqual(dict(net.observed), {})
self.assertEqual(net._original_observed, {})
self.assertEqual(net._stochastic_tensors, {})
with pytest.raises(Exception):
# `net.observed` should be read-only
net.observed['x'] = T.zeros([])
# with observation
normal = UnitNormal([2, 3, 4])
x = T.as_tensor(np.random.randn(3, 4))
y = normal.sample()
net = BayesianNet({'x': x, 'y': y})
self.assertEqual(len(net), 0)
self.assertEqual(list(net), [])
self.assertEqual(list(net.observed), ['x', 'y'])
self.assertIs(net.observed['x'], x)
self.assertIs(net.observed['y'], y.tensor)
self.assertIs(net._original_observed['x'], x)
self.assertIs(net._original_observed['y'], y)
def test_normalize_adj(self):
def D(t):
return np.diag(1. / t)
node_count = 50
eps = 1e-6
# directed
def G(d, y):
return np.dot(D(d), y)
empty_adj = T.sparse.from_dense(T.zeros([node_count, node_count]))
self.assertEqual(
T.shape(T.sparse.get_indices(empty_adj, coord_first=True))[1], 0)
x_list = ([make_random_adj_matrix(node_count)
for _ in range(3)] + [empty_adj])
y_list = [T.sparse.to_numpy(x) for x in x_list]
d_list = [np.maximum(np.sum(y, axis=-1), eps) for y in y_list]
d_sum = sum(d_list, 0.)
for x, y, d in zip(x_list, y_list, d_list):
assert_allclose(gnn.normalize_adj(x, epsilon=eps),
G(d, y),
atol=1e-4,
rtol=1e-6)
out_list = gnn.normalize_partitioned_adj(x_list, epsilon=eps)
for y, out in zip(y_list, out_list):
assert_allclose(out, G(d_sum, y), atol=1e-4, rtol=1e-6)
# undirected
def G(d, y):
d = D(np.sqrt(d))
return np.dot(np.dot(d, y), d)
x_list = [
make_random_adj_matrix(node_count, directed=False)
for _ in range(3)
]
y_list = [T.sparse.to_numpy(x) for x in x_list]
d_list = [np.maximum(np.sum(y, axis=-1), eps) for y in y_list]
d_sum = sum(d_list, 0.)
for x, y, d in zip(x_list, y_list, d_list):
assert_allclose(gnn.normalize_adj(x, undirected=True, epsilon=eps),
G(d, y),
atol=1e-4,
rtol=1e-6)
out_list = gnn.normalize_partitioned_adj(x_list,
undirected=True,
epsilon=eps)
for y, out in zip(y_list, out_list):
assert_allclose(out, G(d_sum, y), atol=1e-4, rtol=1e-6)
# errors
with pytest.raises(Exception,
match='`adj_matrices` must not be empty'):
_ = gnn.normalize_partitioned_adj([])
def test_IdentityFlow(self):
x = T.random.randn([2, 3, 4, 5])
for event_ndims in (0, 1, 2):
flow = tk.layers.jit_compile(tk.flows.IdentityFlow(event_ndims))
log_det_shape = T.shape(x)[:4 - event_ndims]
expected_log_det = T.zeros(log_det_shape)
flow_standard_check(self, flow, x, x, expected_log_det,
T.random.randn(log_det_shape))
def test_sample(self):
t0 = T.zeros([], T.float32)
d = Distribution(
T.float32, [2, 3], continuous=True, reparameterized=True,
event_ndims=1, min_event_ndims=0)
d._sample = Mock(return_value=t0)
# without group_ndims or reparameterized
t = d.sample()
self.assertIs(t, t0)
self.assertEqual(d._sample.call_args, ((None, 0, 1, True), {}))
t = d.sample(n_samples=123)
self.assertEqual(d._sample.call_args, ((123, 0, 1, True), {}))
# with group_ndims
t = d.sample(group_ndims=1)
self.assertEqual(d._sample.call_args, ((None, 1, 2, True), {}))
t = d.sample(n_samples=123, group_ndims=1)
self.assertEqual(d._sample.call_args, ((123, 1, 2, True), {}))
t = d.sample(n_samples=123, group_ndims=2)
self.assertEqual(d._sample.call_args, ((123, 2, 3, True), {}))
with pytest.raises(Exception,
match=re.compile('`min_event_ndims - event_ndims <= group_ndims <= .*'
'sample_ndims - event_ndims` does not hold: ',
re.DOTALL)):
_ = d.sample(group_ndims=2)
with pytest.raises(Exception,
match=re.compile('`min_event_ndims - event_ndims <= group_ndims <= .*'
'sample_ndims - event_ndims` does not hold: ',
re.DOTALL)):
_ = d.sample(n_samples=123, group_ndims=3)
d._sample.reset_mock()
# with reparameterized
t = d.sample(reparameterized=False)
self.assertEqual(d._sample.call_args, ((None, 0, 1, False), {}))
d._sample.reset_mock()
d.reparameterized = False
t = d.sample(reparameterized=False)
self.assertEqual(d._sample.call_args, ((None, 0, 1, False), {}))
d._sample.reset_mock()
with pytest.raises(ValueError,
match='Distribution .* is not re-parameterizable, '
'thus `reparameterized` cannot be set to True'):
_ = d.sample(reparameterized=True)
def check_reversing_flow(ctx, spatial_ndims: int, cls):
num_features = 5
for batch_shape in ([2], [2, 3]):
shape = make_conv_shape(batch_shape, num_features, [6, 7,
8][:spatial_ndims])
flow = tk.layers.jit_compile(cls())
# prepare for the answer
x = T.random.randn(shape)
channel_axis = get_channel_axis(spatial_ndims)
perm = T.arange(num_features - 1, -1, -1, dtype=T.index_dtype)
expected_y = T.index_select(x, perm, axis=channel_axis)
expected_log_det = T.zeros(batch_shape)
# check the flow
flow_standard_check(ctx, flow, x, expected_y, expected_log_det,
T.random.randn(batch_shape))
def test_space_depth_transform(self):
for spatial_ndims, batch_shape, block_size in product(
(1, 2, 3),
([2], [2, 3]),
(1, 2, 4),
):
# prepare for the data
n_channels = 5
x = T.random.randn(
make_conv_shape(batch_shape, n_channels, [4, 8,
12][:spatial_ndims]))
y = getattr(T.nn, f'space_to_depth{spatial_ndims}d')(x, block_size)
log_det = T.zeros(batch_shape)
input_log_det = T.random.randn(batch_shape)
# construct the classes
cls = getattr(tk.flows, f'SpaceToDepth{spatial_ndims}d')
inv_cls = getattr(tk.flows, f'DepthToSpace{spatial_ndims}d')
flow = cls(block_size)
self.assertEqual(flow.block_size, block_size)
inv_flow = inv_cls(block_size)
self.assertEqual(inv_flow.block_size, block_size)
self.assertIsInstance(flow.invert(), inv_cls)
self.assertEqual(flow.invert().block_size, block_size)
self.assertIsInstance(inv_flow.invert(), cls)
self.assertEqual(inv_flow.invert().block_size, block_size)
# check call
flow_standard_check(self, flow, x, y, log_det, input_log_det)
flow_standard_check(self, inv_flow, y, x, log_det, input_log_det)
# test error
with pytest.raises(ValueError,
match='`block_size` must be at least 1'):
_ = cls(0)
with pytest.raises(ValueError,
match='`block_size` must be at least 1'):
_ = inv_cls(0)
def test_deconv(self):
for spatial_ndims in (1, 2, 3):
output_size = [16, 17, 18][:spatial_ndims]
output_shape = make_conv_shape([], 4, output_size)
layer0 = getattr(tk.layers, f'LinearConv{spatial_ndims}d')(
4,
5,
kernel_size=3,
stride=2,
padding='half',
weight_init=tk.init.ones)
y = layer0(T.zeros([1] + output_shape))
input_shape = T.shape(y)[1:]
input_channel, input_size = T.utils.split_channel_spatial_shape(
input_shape)
for fn_name, layer_cls in zip([
f'linear_conv_transpose{spatial_ndims}d',
f'linear_deconv{spatial_ndims}d',
f'conv_transpose{spatial_ndims}d',
f'deconv{spatial_ndims}d',
f'res_block_transpose{spatial_ndims}d'
], [
getattr(tk.layers, f'LinearConvTranspose{spatial_ndims}d'),
getattr(tk.layers, f'LinearConvTranspose{spatial_ndims}d'),
getattr(tk.layers, f'ConvTranspose{spatial_ndims}d'),
getattr(tk.layers, f'ConvTranspose{spatial_ndims}d'),
getattr(tk.layers, f'ResBlockTranspose{spatial_ndims}d'),
]):
# without output_shape
kwargs = {
'kernel_size': 3,
'stride': 2,
'padding': 'half',
'weight_norm': True
}
input_mask = [i == 5 for i in input_shape]
if not fn_name.startswith('linear_'):
kwargs['activation'] = LeakyReLU
sequential_builder_standard_check(ctx=self,
fn_name=fn_name,
layer_cls=layer_cls,
input_shape=input_shape,
input_mask=input_mask,
args=(5, 4),
builder_args=(4, ),
kwargs=kwargs)
kwargs['output_padding'] = 0
sequential_builder_standard_check(ctx=self,
fn_name=fn_name,
layer_cls=layer_cls,
input_shape=input_shape,
input_mask=input_mask,
args=(5, 4),
builder_args=(4, ),
kwargs=kwargs)
# with output_shape
kwargs = {
'kernel_size': 3,
'stride': 2,
'padding': 'half',
'weight_norm': True
}
layer_kwargs = {
'output_padding':
T.utils.calculate_deconv_output_padding(
input_size=input_size,
output_size=output_size,
kernel_size=[3] * spatial_ndims,
stride=[2] * spatial_ndims,
padding=[(1, 1)] * spatial_ndims,
dilation=[1] * spatial_ndims,
)
}
builder_kwargs = {'output_size': output_size}
input_mask = [True] * spatial_ndims
if not fn_name.startswith('linear_'):
kwargs['activation'] = LeakyReLU
sequential_builder_standard_check(
ctx=self,
fn_name=fn_name,
layer_cls=layer_cls,
input_shape=input_shape,
input_mask=input_mask,
args=(5, 4),
builder_args=(4, ),
kwargs=kwargs,
layer_kwargs=layer_kwargs,
builder_kwargs=builder_kwargs,
)
# test errors
builder = SequentialBuilder(input_shape)
fn = getattr(builder, fn_name)
with pytest.raises(ValueError,
match='`output_padding` and `output_size` '
'cannot be both specified'):
fn(5,
kernel_size=1,
output_padding=1,
output_size=[2, 3, 4][:spatial_ndims])
with pytest.raises(ValueError,
match=f'`output_size` is expected to be '
f'{spatial_ndims}d'):
fn(5,
kernel_size=1,
output_size=[2, 3, 4, 5][:spatial_ndims + 1])
builder = SequentialBuilder(
[i if i == 5 else None for i in input_shape])
fn = getattr(builder, fn_name)
with pytest.raises(ValueError,
match='Specifying `output_size` instead of '
'`output_padding` is supported only '
'when the previous output shape '
'is all deterministic.'):
fn(5, kernel_size=1, output_size=[2, 3, 4][:spatial_ndims])
def test_causality_and_receptive_field(self):
for size in [[12], [12, 11], [12, 11, 10]]:
spatial_ndims = len(size)
for kernel_size in [3, 5, [5, 3, 5][:spatial_ndims]]:
# ---- construct the layers ----
# the input layer
input_layer_cls = getattr(tk.layers,
f'PixelCNNInput{spatial_ndims}d')
input_layer = input_layer_cls(
1,
1,
kernel_size=kernel_size,
edge_bias=False,
weight_init=tk.init.ones,
)
input_layer = tk.layers.jit_compile(input_layer)
with pytest.raises(Exception,
match='`input` is expected to be .*d'):
_ = input_layer(T.zeros([1] * (spatial_ndims + 1)))
with pytest.raises(Exception,
match='`input` is expected to be .*d'):
_ = input_layer(T.zeros([1] * (spatial_ndims + 3)))
# `add_ones_channnel = True`
input_layer2 = input_layer_cls(1,
1,
kernel_size=kernel_size,
weight_init=tk.init.ones)
# the pixelcnn resblock
resblock_layer_cls = getattr(
tk.layers, f'PixelCNNResBlock{spatial_ndims}d')
with pytest.raises(ValueError,
match=r'`kernel_size` is required to be at '
r'least 3'):
_ = resblock_layer_cls(1, 1, kernel_size=1)
with pytest.raises(
ValueError,
match=r'`kernel_size` is required to be odd'):
_ = resblock_layer_cls(1,
1,
kernel_size=[4, 3,
5][:spatial_ndims])
resblock_layer = resblock_layer_cls(1,
1,
kernel_size=kernel_size,
weight_init=tk.init.ones)
resblock_layer = tk.layers.jit_compile(resblock_layer)
with pytest.raises(Exception):
_ = resblock_layer([T.zeros([])] * (spatial_ndims - 1))
with pytest.raises(Exception):
_ = resblock_layer([T.zeros([])] * (spatial_ndims + 1))
# the down-sampling and up-sampling layer
down_sample_cls = getattr(tk.layers,
f'PixelCNNConv{spatial_ndims}d')
down_sample_layer = down_sample_cls(1,
1,
kernel_size,
stride=2)
down_sample_layer = tk.layers.jit_compile(down_sample_layer)
down_sample_output_size = T.shape(
down_sample_layer(
[T.zeros(make_conv_shape([1], 1, size))] *
spatial_ndims)[0])
up_sample_cls = getattr(
tk.layers, f'PixelCNNConvTranspose{spatial_ndims}d')
up_sample_layer = up_sample_cls(
1,
1,
kernel_size,
stride=2,
output_padding=tk.layers.get_deconv_output_padding(
input_size=[
down_sample_output_size[a]
for a in get_spatial_axis(spatial_ndims)
],
output_size=size,
kernel_size=kernel_size,
stride=2,
padding=
'half', # sum of the both sides == (kernel_size - 1) * dilation
))
up_sample_layer = tk.layers.jit_compile(up_sample_layer)
# the output layer
output_layer_cls = getattr(tk.layers,
f'PixelCNNOutput{spatial_ndims}d')
output_layer = output_layer_cls()
output_layer = tk.layers.jit_compile(output_layer)
with pytest.raises(
Exception,
match=r'`len\(inputs\)` is expected to be .*'):
_ = output_layer([T.zeros([])] * (spatial_ndims - 1))
with pytest.raises(
Exception,
match=r'`len\(inputs\)` is expected to be .*'):
_ = output_layer([T.zeros([])] * (spatial_ndims + 1))
# ---- test the causality ----
for pos, single_point in product(iter_causal_test_pos(size),
(True, False)):
x = make_causal_test_input(size,
pos,
single_point=single_point)
x_t = T.as_tensor(x)
# check the input layer output
outputs = input_layer(x_t)
ensure_stacks_causality(self, outputs, size, pos)
# check the final output
assert_allclose(output_layer(outputs), outputs[-1])
# check the resblock output
resblock_outputs = resblock_layer(outputs)
ensure_stacks_causality(self, resblock_outputs, size, pos)
outputs2 = resblock_outputs
for i in range(4):
outputs2 = resblock_layer(outputs2)
ensure_full_receptive_field(self, outputs2[-1], size, pos)
# check the down-sample and up-sample
down_sample_outputs = down_sample_layer(outputs)
up_sample_outputs = up_sample_layer(down_sample_outputs)
ensure_stacks_causality(self, up_sample_outputs, size, pos)
# ---- test zero input on different input layers ----
x_t = T.zeros(make_conv_shape([1], 1, size), dtype=T.float32)
outputs = input_layer(x_t)
assert_equal(
(np.abs(T.to_numpy(outputs[-1])) >= 1e-6).astype(np.int32),
x_t)
outputs = input_layer2(x_t)
assert_equal(
(np.abs(T.to_numpy(outputs[-1])) >= 1e-6).astype(np.int32),
make_causal_mask(size,
[0] * spatial_ndims).astype(np.int32))
def test_conv_deconv_output_shape_and_args(self):
for input_size, kernel_size, stride, padding, dilation in product(
([8, 9, 10], [16, 21, 32], [30, 31, 32]),
([1] * 3, [2] * 3, [3] * 3, [1, 2, 3]),
([1] * 3, [2] * 3, [3] * 3, [1, 2, 3]),
([(0, 0)] * 3, [(1, 1)] * 3, [(2, 2)] * 3, [(3, 3)] * 3, [(1, 2),
(2, 3),
(3, 4)]),
([1] * 3, [2] * 3, [3] * 3, [1, 2, 3]),
):
args = (input_size, kernel_size, stride, padding, dilation)
# calculate_conv_output_size
output_size = [get_conv_output_size(*a) for a in zip(*args)]
self.assertEqual(
T.utils.calculate_conv_output_size(
input_size=input_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
), output_size)
layer1 = tk.layers.LinearConv3d(
1,
1,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
)
x = T.zeros(make_conv_shape([1], 1, input_size))
y = layer1(x)
self.assertEqual(
T.utils.split_channel_spatial_shape(T.shape(y)[1:])[1],
output_size,
)
# calculate_deconv_output_padding
output_padding = T.utils.calculate_deconv_output_padding(
input_size=output_size,
output_size=input_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
)
layer2 = tk.layers.LinearConvTranspose3d(
1,
1,
kernel_size=kernel_size,
stride=stride,
padding=padding,
output_padding=output_padding,
dilation=dilation,
)
z = layer2(y)
self.assertEqual(
T.utils.split_channel_spatial_shape(T.shape(z)[1:])[1],
input_size,
)
# calculate_deconv_output_size
self.assertEqual(
T.utils.calculate_deconv_output_size(
input_size=output_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
output_padding=output_padding,
dilation=dilation,
), input_size)
# test error
kwargs = dict(kernel_size=[1],
stride=[1],
dilation=[1],
padding=[(0, 0)])
for input_size in ([], [1, 2, 3, 4]):
with pytest.raises(Exception,
match='`input_size` is not a 1d, 2d or 3d '
'convolutional input size'):
_ = T.utils.calculate_conv_output_size(input_size, **kwargs)
with pytest.raises(Exception,
match='`input_size` is not a 1d, 2d or 3d '
'convolutional input size'):
_ = T.utils.calculate_deconv_output_size(input_size,
output_padding=[0],
**kwargs)
for arg_name in ('kernel_size', 'stride', 'dilation', 'padding'):
kwargs2 = dict(kwargs)
if arg_name == 'padding':
kwargs2[arg_name] = [(0, 0)] * 2
else:
kwargs2[arg_name] = [1, 1]
with pytest.raises(Exception,
match='`.*` is not for .*d convolution'):
_ = T.utils.calculate_conv_output_size([11], **kwargs2)
with pytest.raises(Exception,
match='`.*` is not for .*d convolution'):
_ = T.utils.calculate_deconv_output_size([11],
output_padding=[0],
**kwargs2)
with pytest.raises(Exception, match='`.*` is not for .*d convolution'):
_ = T.utils.calculate_deconv_output_size([11],
output_padding=[0, 0],
**kwargs)
with pytest.raises(Exception,
match='No `output_padding` can satisfy the '
'deconvolution task'):
_ = T.utils.calculate_deconv_output_padding([2], [1], [1], [1],
[(0, 0)], [1])
def test_construct_base(self):
logits = np.random.randn(2, 3, 4)
probs = softmax(logits)
logits = np.log(probs)
for dtype, float_dtype in product(number_dtypes, float_dtypes):
logits_t = T.as_tensor(logits, dtype=float_dtype)
probs_t = T.as_tensor(probs, dtype=float_dtype)
mutual_params = {'logits': logits_t, 'probs': probs_t}
# construct from logits or probs
for key, val in mutual_params.items():
other_key = [k for k in mutual_params if k != key][0]
cat = _MyBaseCategorical(event_ndims=1,
dtype=dtype,
epsilon=1e-6,
**{key: val})
self.assertEqual(cat.continuous, False)
self.assertEqual(cat.reparameterized, False)
self.assertEqual(cat.dtype, dtype)
self.assertEqual(cat.event_ndims, 1)
self.assertEqual(cat.epsilon, 1e-6)
self.assertIs(getattr(cat, key), val)
assert_allclose(getattr(cat, other_key),
mutual_params[other_key],
rtol=1e-4)
self.assertEqual(cat._mutual_params, {key: val})
# must specify either logits or probs, but not both
with pytest.raises(ValueError,
match='Either `logits` or `probs` must be '
'specified, but not both.'):
_ = _MyBaseCategorical(logits=logits_t,
probs=probs_t,
dtype=dtype,
event_ndims=0)
with pytest.raises(ValueError,
match='Either `logits` or `probs` must be '
'specified, but not both.'):
_ = _MyBaseCategorical(logits=None,
probs=None,
dtype=dtype,
event_ndims=0)
# shape test on logits or probs
for key in mutual_params:
param_val = T.zeros([], dtype=float_dtype)
with pytest.raises(ValueError,
match=rf'`{key}` must be at least 1d: '
rf'got shape \[\]'):
_ = _MyBaseCategorical(dtype=dtype,
event_ndims=0,
**{key: param_val})
# nan test
for key, val in mutual_params.items():
with pytest.raises(Exception,
match='Infinity or NaN value encountered'):
_ = _MyBaseCategorical(validate_tensors=True,
dtype=dtype,
event_ndims=2,
**{
key:
T.as_tensor(np.asarray(
[[np.nan]]),
dtype=float_dtype)
})
def test_ctor_and_type_convert(self):
def f(row, col, values, shape, dtype, force_coalesced):
values = values.astype(np.float64)
if dtype == T.int32 or dtype == T.int64:
values = np.asarray(values * 1000, dtype=np.int64)
x = make_ndarray_by_coo(row, col, values, shape)
self.assertFalse(T.sparse.is_sparse_tensor(x))
self.assertFalse(T.sparse.is_sparse_tensor(T.as_tensor(x)))
the_force_coalesced = (
force_coalesced if force_coalesced is not None else
T.sparse.MAKE_SPARSE_DEFAULT_FORCE_COALESCED)
the_dtype = dtype if dtype is not None else T.float64
def g(x, y):
self.assertTrue(T.sparse.is_sparse_tensor(y))
self.assertEqual(T.sparse.is_coalesced(y), the_force_coalesced)
self.assertEqual(T.sparse.get_dtype(y), the_dtype)
self.assertEqual(T.sparse.rank(y), len(shape))
self.assertEqual(T.sparse.length(y), shape[0])
self.assertEqual(T.sparse.shape(y), shape)
self.assertEqual(T.sparse.get_device(y), T.current_device())
assert_allclose(x, y, rtol=1e-4, atol=1e-6)
def h(x, y, t_):
self.assertIsInstance(y, t_)
assert_allclose(x, y)
# make_sparse
for coord_first in (None, True, False):
kwargs = {}
if coord_first is not None:
kwargs['coord_first'] = coord_first
if force_coalesced is not None:
kwargs['force_coalesced'] = force_coalesced
the_coord_first = (
coord_first if coord_first is not None else
T.sparse.SPARSE_INDICES_DEFAULT_IS_COORD_FIRST)
y = T.sparse.make_sparse(T.as_tensor(np.stack(
[row, col], axis=0 if the_coord_first else 1),
dtype=T.int64),
T.as_tensor(values),
dtype=dtype,
shape=shape,
**kwargs)
g(x, y)
self.assertEqual(T.sparse.value_count(y), len(row))
# from_dense
y = T.sparse.from_dense(
T.as_tensor(x, dtype=dtype),
**({
'force_coalesced': force_coalesced
} if force_coalesced is not None else {}))
self.assertTrue(T.sparse.is_sparse_tensor(y))
g(x, y)
t = T.sparse.to_dense(y)
self.assertFalse(T.sparse.is_sparse_tensor(t))
h(x, t, T.Tensor)
# from_numpy
y = T.sparse.from_numpy(x,
dtype=dtype,
**({
'force_coalesced': force_coalesced
} if force_coalesced is not None else {}))
g(x, y)
h(x, T.sparse.to_numpy(y), np.ndarray)
# from_coomatrix
if len(shape) == 2:
spmat = sp.coo_matrix((values, (row, col)), shape=shape)
for m in (spmat, spmat.tocsr()):
y = T.sparse.from_spmatrix(
m,
dtype=dtype,
**({
'force_coalesced': force_coalesced
} if force_coalesced is not None else {}))
g(x, y)
h(m, T.sparse.to_spmatrix(y), sp.spmatrix)
# dtype from another sparse tensor
z = T.sparse.make_sparse(
T.as_tensor(np.stack([row, col], axis=0)),
T.as_tensor(values),
dtype=y.dtype,
shape=shape,
)
self.assertEqual(T.sparse.get_dtype(z), T.sparse.get_dtype(y))
# test to_dtype
z = T.sparse.make_sparse(
T.as_tensor(np.stack([row, col], axis=0)),
T.as_tensor(values),
shape=shape,
)
z = T.sparse.to_dtype(z, T.sparse.get_dtype(y))
self.assertEqual(T.sparse.get_dtype(z), T.sparse.get_dtype(y))
# test ordinary
for force_coalesced in [None, True, False]:
for dtype in [None, 'float32', T.float64, T.int32, 'int64']:
f(row=np.array([0, 0, 1, 3, 4]),
col=np.array([1, 4, 5, 3, 2]),
values=np.random.randn(5),
shape=[5, 6],
dtype=dtype,
force_coalesced=force_coalesced)
f(row=np.array([0, 0, 1, 3, 4]),
col=np.array([1, 4, 5, 3, 2]),
values=np.random.randn(10).reshape([5, 2]),
shape=[5, 6, 2],
dtype=dtype,
force_coalesced=force_coalesced)
# test with_device and to_device
f = lambda: T.sparse.make_sparse(
T.as_tensor([[0, 1], [1, 0]]), T.random.randn([2]), shape=[3, 3])
with T.use_device(T.CPU_DEVICE):
t = f()
self.assertEqual(T.sparse.get_device(t), T.CPU_DEVICE)
t = T.sparse.to_device(f(), T.CPU_DEVICE)
self.assertEqual(T.sparse.get_device(t), T.CPU_DEVICE)
# test errors
with pytest.raises(ValueError, match='`indices` must be a 2d tensor'):
_ = T.sparse.make_sparse(T.zeros([2, 3, 4]),
T.random.randn([5]),
shape=[5, 5])
with pytest.raises(ValueError, match='`dtype` not supported'):
_ = T.sparse.make_sparse(
T.as_tensor([[0, 1], [1, 0]]),
T.random.randn([2]),
shape=[3, 3],
dtype=T.boolean,
)
with pytest.raises(ValueError,
match='`indices` must be a int32 or '
'int64 tensor'):
_ = T.sparse.make_sparse(T.as_tensor([[0, 1], [1, 0]],
dtype=T.int16),
T.random.randn([2]),
shape=[3, 3])
def do_check(batch_shape, scale_type, initialized, dtype):
x = T.random.randn(make_conv_shape(batch_shape, num_features,
[6, 7, 8][:spatial_ndims]),
dtype=dtype)
# check construct
flow = cls(num_features,
scale=scale_type,
initialized=initialized,
dtype=dtype)
ctx.assertIn(f'num_features={num_features}', repr(flow))
ctx.assertIn(f'axis={-(spatial_ndims + 1)}', repr(flow))
ctx.assertIn(f'scale_type={scale_type!r}', repr(flow))
flow = tk.layers.jit_compile(flow)
# check initialize
if not initialized:
# must initialize with sufficient data
with pytest.raises(Exception, match='with at least .* dimensions'):
_ = flow(
T.random.randn(make_conv_shape([], num_features,
[6, 7, 8][:spatial_ndims]),
dtype=dtype))
# must initialize with inverse = Fale
with pytest.raises(Exception,
match='`ActNorm` must be initialized with '
'`inverse = False`'):
_ = flow(x, inverse=True)
# do initialize
y, _ = flow(x, compute_log_det=False)
y_mean, y_var = T.calculate_mean_and_var(
y, axis=[a for a in range(-T.rank(y), 0) if a != channel_axis])
assert_allclose(y_mean,
T.zeros([num_features]),
rtol=1e-4,
atol=1e-6)
assert_allclose(y_var,
T.ones([num_features]),
rtol=1e-4,
atol=1e-6)
else:
y, _ = flow(x, compute_log_det=False)
assert_allclose(y, x, rtol=1e-4, atol=1e-6)
# prepare for the expected result
scale_obj = ExpScale() if scale_type == 'exp' else LinearScale()
if T.IS_CHANNEL_LAST:
aligned_shape = [num_features]
else:
aligned_shape = [num_features] + [1] * spatial_ndims
bias = T.reshape(flow.bias, aligned_shape)
pre_scale = T.reshape(flow.pre_scale, aligned_shape)
expected_y, expected_log_det = scale_obj(x + bias,
pre_scale,
event_ndims=(spatial_ndims +
1),
compute_log_det=True)
flow_standard_check(ctx, flow, x, expected_y, expected_log_det,
T.random.randn(T.shape(expected_log_det)))
def test_discretized_logistic(self):
next_seed_val = [1234]
def next_seed():
ret = next_seed_val[0]
next_seed_val[0] += 1
return ret
def safe_sigmoid(x):
return np.where(x < 0,
np.exp(x) / (1. + np.exp(x)),
1. / (1. + np.exp(-x)))
def do_discretize(x, bin_size, min_val=None, max_val=None):
if min_val is not None:
x = x - min_val
x = np.floor(x / bin_size + .5) * bin_size
if min_val is not None:
x = x + min_val
if min_val is not None:
x = np.maximum(x, min_val)
if max_val is not None:
x = np.minimum(x, max_val)
return x
def naive_discretized_logistic_pdf(x,
mean,
log_scale,
bin_size,
min_val=None,
max_val=None,
biased_edges=True,
discretize_given=True):
# discretize x
if discretize_given:
x = do_discretize(x, bin_size, min_val, max_val)
# middle pdfs
half_bin = bin_size * 0.5
x_hi = (x - mean + half_bin) / np.exp(log_scale)
x_low = (x - mean - half_bin) / np.exp(log_scale)
cdf_delta = safe_sigmoid(x_hi) - safe_sigmoid(x_low)
middle_pdf = np.log(np.maximum(cdf_delta, 1e-7))
log_prob = middle_pdf
# left edge
if min_val is not None and biased_edges:
log_prob = np.where(x < min_val + half_bin,
np.log(safe_sigmoid(x_hi)), log_prob)
# right edge
if max_val is not None and biased_edges:
log_prob = np.where(x >= max_val - half_bin,
np.log(1. - safe_sigmoid(x_low)), log_prob)
# zero out prob outside of [min_val - half_bin, max_val + half_bin].
if min_val is not None and max_val is not None:
log_prob = np.where(
np.logical_and(
x >= min_val - half_bin,
x <= max_val + half_bin,
), log_prob, T.random.LOG_ZERO_VALUE)
return log_prob
def naive_discretized_logistic_sample(uniform_samples,
mean,
log_scale,
bin_size,
min_val=None,
max_val=None,
discretize_sample=True):
u = uniform_samples
samples = mean + np.exp(log_scale) * (np.log(u) - np.log1p(-u))
if discretize_sample:
samples = do_discretize(samples, bin_size, min_val, max_val)
return samples
def get_samples(mean, log_scale, n_samples=None, **kwargs):
seed = next_seed()
kwargs.setdefault('epsilon', 1e-7)
sample_shape = T.get_broadcast_shape(T.shape(mean),
T.shape(log_scale))
if n_samples is not None:
sample_shape = [n_samples] + sample_shape
np.random.seed(seed)
T.random.seed(seed)
u = T.random.uniform(shape=sample_shape,
low=kwargs['epsilon'],
high=1. - kwargs['epsilon'],
dtype=T.get_dtype(mean))
u = T.to_numpy(u)
np.random.seed(seed)
T.random.seed(seed)
r = T.random.discretized_logistic(mean,
log_scale,
n_samples=n_samples,
**kwargs)
return u, r
mean = 3 * np.random.uniform(size=[2, 1, 4]).astype(np.float64) - 1
log_scale = np.random.normal(size=[3, 1, 5, 1]).astype(np.float64)
# sample
def do_test_sample(bin_size: float, min_val: Optional[float],
max_val: Optional[float], discretize_sample: bool,
discretize_given: bool, biased_edges: bool,
reparameterized: bool, n_samples: Optional[int],
validate_tensors: bool, dtype: str):
mean_t = T.as_tensor(mean, dtype=dtype)
log_scale_t = T.as_tensor(log_scale, dtype=dtype)
value_shape = T.get_broadcast_shape(T.shape(mean_t),
T.shape(log_scale_t))
# sample
sample_shape = [n_samples] if n_samples is not None else []
u, t = get_samples(
mean_t,
log_scale_t,
n_samples=n_samples,
bin_size=bin_size,
min_val=min_val,
max_val=max_val,
discretize=discretize_sample,
reparameterized=reparameterized,
epsilon=T.EPSILON,
validate_tensors=validate_tensors,
)
self.assertEqual(T.get_dtype(t), dtype)
self.assertEqual(T.get_device(t), T.current_device())
self.assertEqual(T.shape(t), sample_shape + value_shape)
# check values
this_mean = mean.astype(dtype)
this_log_scale = log_scale.astype(dtype)
expected_t = naive_discretized_logistic_sample(
u,
this_mean,
this_log_scale,
bin_size,
min_val,
max_val,
discretize_sample=discretize_sample,
)
assert_allclose(t, expected_t, rtol=1e-4, atol=1e-6)
# check log_prob
do_check_log_prob(given=t,
batch_ndims=len(t.shape),
Z_log_prob_fn=partial(
T.random.discretized_logistic_log_prob,
mean=mean_t,
log_scale=log_scale_t,
bin_size=bin_size,
min_val=min_val,
max_val=max_val,
biased_edges=biased_edges,
discretize=discretize_given,
validate_tensors=validate_tensors,
),
np_log_prob=naive_discretized_logistic_pdf(
x=T.to_numpy(t),
mean=this_mean,
log_scale=this_log_scale,
bin_size=bin_size,
min_val=min_val,
max_val=max_val,
biased_edges=biased_edges,
discretize_given=discretize_given,
))
for dtype in float_dtypes:
do_test_sample(bin_size=1 / 31.,
min_val=None,
max_val=None,
discretize_sample=True,
discretize_given=True,
biased_edges=False,
reparameterized=False,
n_samples=None,
validate_tensors=False,
dtype=dtype)
do_test_sample(bin_size=1 / 32.,
min_val=-3.,
max_val=2.,
discretize_sample=True,
discretize_given=True,
biased_edges=True,
reparameterized=False,
n_samples=100,
validate_tensors=True,
dtype=dtype)
for discretize, biased_edges, validate_tensors in product(
[True, False],
[True, False],
[True, False],
):
do_test_sample(bin_size=1 / 127.,
min_val=None,
max_val=None,
discretize_sample=discretize,
discretize_given=discretize,
biased_edges=biased_edges,
reparameterized=not discretize,
n_samples=None,
validate_tensors=validate_tensors,
dtype=T.float32)
do_test_sample(bin_size=1 / 128.,
min_val=-3.,
max_val=2.,
discretize_sample=discretize,
discretize_given=discretize,
biased_edges=biased_edges,
reparameterized=not discretize,
n_samples=None,
validate_tensors=validate_tensors,
dtype=T.float32)
mean_t = T.as_tensor(mean, dtype=T.float32)
log_scale_t = T.as_tensor(log_scale, dtype=T.float32)
given_t = T.zeros(
T.get_broadcast_shape(T.shape(mean_t), T.shape(log_scale_t)))
with pytest.raises(
Exception,
match=
'`min_val` and `max_val` must be both None or neither None'):
_ = T.random.discretized_logistic(mean_t,
log_scale_t,
1 / 255.,
min_val=-3.)
with pytest.raises(
Exception,
match=
'`min_val` and `max_val` must be both None or neither None'):
_ = T.random.discretized_logistic(mean_t,
log_scale_t,
1 / 255.,
max_val=2.)
with pytest.raises(
Exception,
match='`discretize` and `reparameterized` cannot be both True'
):
_ = T.random.discretized_logistic(mean_t,
log_scale_t,
1 / 255.,
discretize=True,
reparameterized=True)
with pytest.raises(Exception,
match='`mean.dtype` != `log_scale.dtype`'):
_ = T.random.discretized_logistic(
mean_t, T.as_tensor(log_scale, dtype=T.float64), 1 / 255.)
with pytest.raises(
Exception,
match=
'`min_val` and `max_val` must be both None or neither None'):
_ = T.random.discretized_logistic_log_prob(given_t,
mean_t,
log_scale_t,
1 / 255.,
min_val=-3.)
with pytest.raises(
Exception,
match=
'`min_val` and `max_val` must be both None or neither None'):
_ = T.random.discretized_logistic_log_prob(given_t,
mean_t,
log_scale_t,
1 / 255.,
max_val=2.)
def net_builder(observed):
net = BayesianNet(observed)
z = net.add('z', UnitNormal([1]))
y = net.add('y', Normal(T.zeros([1]), T.full([1], 2.)))
x = net.add('x', Normal(z.tensor + y.tensor, T.ones([1])))
return net
def test_discretized_logsitic(self):
mean = T.random.randn([3, 1, 4])
log_scale = T.random.randn([2, 1])
def do_check(**kwargs):
d = DiscretizedLogistic(mean, log_scale, **kwargs)
event_ndims = kwargs.get('event_ndims', 0)
value_shape = T.get_broadcast_shape(T.shape(mean), T.shape(log_scale))
log_prob_fn_kwargs = copy.copy(kwargs)
log_prob_fn_kwargs.pop('discretize_sample', None)
log_prob_fn_kwargs['discretize'] = \
log_prob_fn_kwargs.pop('discretize_given', True)
def log_prob_fn(t):
return T.random.discretized_logistic_log_prob(
T.as_tensor(t), mean=mean, log_scale=log_scale,
**log_prob_fn_kwargs
)
check_distribution_instance(
ctx=self,
d=d,
event_ndims=event_ndims,
batch_shape=value_shape[: len(value_shape) - event_ndims],
event_shape=value_shape[len(value_shape) - event_ndims:],
min_event_ndims=0,
max_event_ndims=len(value_shape),
log_prob_fn=log_prob_fn,
# other attributes,
**kwargs
)
for biased_edges, discretize_given, discretize_sample in product(
[True, False],
[True, False],
[True, False],
):
do_check(bin_size=1. / 255,
biased_edges=biased_edges,
discretize_given=discretize_given,
discretize_sample=discretize_sample)
do_check(bin_size=1. / 255,
biased_edges=biased_edges,
discretize_given=discretize_given,
discretize_sample=discretize_sample,
min_val=-3.,
max_val=2.)
with pytest.raises(ValueError,
match='`reparameterized` cannot be True when '
'`discretize_sample` is True'):
_ = DiscretizedLogistic(mean, log_scale, 1./32, reparameterized=True,
discretize_sample=True)
with pytest.raises(ValueError,
match='`min_val - max_val` must be multiples of '
'`bin_size`'):
_ = DiscretizedLogistic(mean, log_scale, .3, min_val=.1, max_val=.2)
with pytest.raises(ValueError,
match='`min_val` and `max_val` must be both None or '
'neither None'):
_ = DiscretizedLogistic(mean, log_scale, 1./32, max_val=2.)
with pytest.raises(ValueError,
match='`min_val` and `max_val` must be both None or '
'neither None'):
_ = DiscretizedLogistic(mean, log_scale, 1./32, min_val=-3.)
with pytest.raises(ValueError,
match='The shape of `mean` and `log_scale` cannot be '
'broadcasted against each other'):
_ = DiscretizedLogistic(mean, T.zeros([7]), 1./32)
