def test_InverseFlow(self): original_flow = tk.layers.jit_compile(_MyFlow()) flow = InverseFlow(original_flow) self.assertIs(flow.original_flow, original_flow) self.assertIs(flow.invert(), original_flow) flow = tk.layers.jit_compile(flow) self.assertEqual(flow.get_x_event_ndims(), 2) self.assertEqual(flow.get_y_event_ndims(), 1) self.assertTrue(flow.is_explicitly_invertible()) x = T.random.randn([2, 3, 4, 1]) expected_y = T.reshape((x - 1.) * 0.5, [2, 3, 4]) expected_log_det = -T.full([2, 3], math.log(2.) * 4) input_log_det = T.random.randn([2, 3]) flow_standard_check(self, flow, x, expected_y, expected_log_det, input_log_det) with pytest.raises(TypeError, match='`flow` must be an explicitly invertible flow'): _ = InverseFlow(tk.layers.Linear(5, 3)) base_flow = _MyFlow() base_flow.explicitly_invertible = False with pytest.raises(TypeError, match='`flow` must be an explicitly invertible flow'): _ = InverseFlow(tk.layers.jit_compile(base_flow))
def test_copy(self): for dtype in float_dtypes: low_t = T.full([2, 1], -1., dtype=dtype) high_t = T.full([1, 3], 2., dtype=dtype) uniform = Uniform(shape=[5, 4], low=low_t, high=high_t, event_ndims=1, log_zero=-1e6, reparameterized=False) self.assertIs(uniform.low, low_t) self.assertIs(uniform.high, high_t) self.assertEqual(uniform.reparameterized, False) self.assertEqual(uniform.value_shape, [5, 4, 2, 3]) self.assertEqual(uniform.dtype, dtype) self.assertEqual(uniform.event_ndims, 1) self.assertEqual(uniform.log_zero, -1e6) with mock.patch( 'tensorkit.distributions.uniform.copy_distribution', wraps=copy_distribution) as f_copy: uniform2 = uniform.copy(event_ndims=2) self.assertIsInstance(uniform2, Uniform) self.assertIs(uniform2.low, low_t) self.assertIs(uniform2.high, high_t) self.assertEqual(uniform2.value_shape, [5, 4, 2, 3]) self.assertEqual(uniform2.dtype, dtype) self.assertEqual(uniform2.event_ndims, 2) self.assertEqual(uniform2.log_zero, -1e6) self.assertEqual(f_copy.call_args, ((), { 'cls': Uniform, 'base': uniform, 'attrs': (('shape', '_shape'), 'low', 'high', 'dtype', 'reparameterized', 'event_ndims', 'log_zero', 'device', 'validate_tensors'), 'overrided_params': { 'event_ndims': 2 }, }))
def test_call(self): # test call and inverse call flows = [_MyFlow1(), tk.layers.jit_compile(_MyFlow1())] flow = tk.layers.jit_compile(SequentialFlow(flows)) x = T.random.randn([2, 3, 4]) expected_y = (x * 2. + 1.) * 2. + 1. expected_log_det = T.full([2, 3], math.log(2.) * 8) input_log_det = T.random.randn([2, 3]) flow_standard_check(self, flow, x, expected_y, expected_log_det, input_log_det) # test no inverse call flows = [_MyFlow1()] flows[0].explicitly_invertible = False flow = tk.layers.jit_compile(SequentialFlow(flows)) with pytest.raises(Exception, match='Flow is not explicitly invertible'): _ = flow(x, inverse=True)
def test_call(self): flow = tk.layers.jit_compile(_MyFlow()) self.assertEqual(flow.get_x_event_ndims(), 1) self.assertEqual(flow.get_y_event_ndims(), 2) self.assertEqual(flow.is_explicitly_invertible(), True) # test call x = T.random.randn([2, 3, 4]) expected_y = T.reshape(x * 2. + 1., [2, 3, 4, 1]) expected_log_det = T.full([2, 3], math.log(2.) * 4) input_log_det = T.random.randn([2, 3]) flow_standard_check(self, flow, x, expected_y, expected_log_det, input_log_det) # test input shape error with pytest.raises(Exception, match='`input` is required to be at least .*d'): _ = flow(T.random.randn([])) with pytest.raises(Exception, match='`input` is required to be at least .*d'): _ = flow(T.random.randn([3]), inverse=True) # test input_log_det shape error with pytest.raises(Exception, match='The shape of `input_log_det` is not expected'): _ = flow(x, T.random.randn([2, 4])) with pytest.raises(Exception, match='The shape of `input_log_det` is not expected'): _ = flow(expected_y, T.random.randn([2, 4]), inverse=True) # test output_log_det shape error flow = tk.layers.jit_compile(_MyBadFlow()) with pytest.raises(Exception, match='(shape|size)'): _ = flow(x) with pytest.raises(Exception, match='(shape|size)'): _ = flow(x, inverse=True)
def test_construct(self): for dtype in float_dtypes: # specify no args uniform = Uniform(dtype=dtype, event_ndims=0) self.assertEqual(uniform.value_shape, []) self.assertEqual(uniform.low, None) self.assertEqual(uniform.high, None) self.assertEqual(uniform.event_ndims, 0) self.assertEqual(uniform.log_zero, -1e7) # specify `low` and `high` float uniform = Uniform(low=-1., high=2., dtype=dtype, event_ndims=0) self.assertEqual(uniform.value_shape, []) self.assertEqual(uniform.low, -1.) self.assertEqual(uniform.high, 2.) self.assertEqual(uniform.event_ndims, 0) # specify `low`, `high` tensors low_t = T.full([2, 1], -1., dtype=dtype) high_t = T.full([1, 3], 2., dtype=dtype) uniform = Uniform(low=low_t, high=high_t, dtype=dtype, event_ndims=2) self.assertEqual(uniform.value_shape, [2, 3]) self.assertEqual(uniform.dtype, dtype) self.assertEqual(uniform.event_ndims, 2) self.assertIs(uniform.low, low_t) self.assertIs(uniform.high, high_t) # specify `low` or `high`, one as tensor and one as numpy array for low, high in [(low_t, T.to_numpy(high_t)), (T.to_numpy(low_t), high_t)]: uniform = Uniform( low=low, high=high, dtype=T.float32, # should be ignored event_ndims=2) self.assertEqual(uniform.value_shape, [2, 3]) self.assertEqual(uniform.dtype, dtype) self.assertEqual(T.get_dtype(uniform.low), dtype) self.assertEqual(T.get_dtype(uniform.high), dtype) self.assertEqual(uniform.event_ndims, 2) assert_equal(uniform.low, low_t) assert_equal(uniform.high, high_t) for event_ndims, dtype, shape in product(range(0, 3), float_dtypes, ([], [2, 3])): if event_ndims > len(shape): continue # specify `shape` uniform = Uniform(shape=shape, dtype=dtype, event_ndims=event_ndims) self.assertEqual(uniform.value_shape, shape) self.assertEqual(uniform.dtype, dtype) self.assertEqual(uniform.event_ndims, event_ndims) self.assertEqual(uniform.low, None) self.assertEqual(uniform.high, None) # specify `shape` and `low`, `high` floats uniform = Uniform(shape=shape, low=-1., high=2., dtype=dtype, event_ndims=event_ndims) self.assertEqual(uniform.value_shape, shape) self.assertEqual(uniform.dtype, dtype) self.assertEqual(uniform.event_ndims, event_ndims) self.assertEqual(uniform.low, -1.) self.assertEqual(uniform.high, 2.) # specify just one of `low`, `high` as float, another as tensor uniform = Uniform(shape=shape, low=-1., high=T.as_tensor(2., dtype=dtype), event_ndims=event_ndims) self.assertEqual(uniform.value_shape, shape) self.assertEqual(uniform.dtype, dtype) self.assertEqual(uniform.event_ndims, event_ndims) self.assertEqual(uniform.low, -1.) self.assertEqual(uniform.high, 2.) for event_ndims, dtype, shape in product(range(0, 3), float_dtypes, ([], [2, 3])): if event_ndims > len(shape) + 2: continue # specify `shape` and `low`, `high` tensors low_t = T.full([2, 1], -1., dtype=dtype) high_t = T.full([1, 3], 2., dtype=dtype) uniform = Uniform( shape=shape, low=low_t, high=high_t, dtype=T.float32, # should be ignored event_ndims=event_ndims) self.assertEqual(uniform.value_shape, shape + [2, 3]) self.assertEqual(uniform.dtype, dtype) self.assertEqual(uniform.event_ndims, event_ndims) self.assertIs(uniform.low, low_t) self.assertIs(uniform.high, high_t) with pytest.raises(ValueError, match='`low` and `high` must be both specified, or ' 'neither specified'): _ = Uniform(low=-1.) with pytest.raises( ValueError, match='`high.dtype` != `low.dtype`: float64 vs float32'): _ = Uniform(low=T.full([2, 3], -1., dtype=T.float32), high=T.full([2, 3], 2., dtype=T.float64)) with pytest.raises(ValueError, match='`low` < `high` does not hold: `low` == 2.0, ' '`high` == 1.0'): _ = Uniform(low=2., high=1.) with pytest.raises(Exception, match='`low` < `high` does not hold'): _ = Uniform(low=T.full([2, 3], 2., dtype=T.float32), high=T.full([2, 3], -1., dtype=T.float32), validate_tensors=True)
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