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