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)