def test_sequential_with_quadratic_flows(self): n_layers = 3 flow1 = MultiLayerQuadraticFlow(n_layers) flow2 = SequentialFlow([ QuadraticFlow(i + 1., i * 2. + 1.) for i in range(n_layers) ]) self.assertTrue(flow2.explicitly_invertible) self.assertEqual(len(flow2.flows), n_layers) for i in range(n_layers): self.assertEqual(flow2.flows[i].a, i + 1.) self.assertEqual(flow2.flows[i].b, i * 2. + 1.) x = tf.range(12, dtype=tf.float32) + 1. with self.test_session() as sess: invertible_flow_standard_check(self, flow2, sess, x) # transform y1, log_det_y1 = flow1.transform(x) y2, log_det_y2 = flow2.transform(x) np.testing.assert_allclose(*sess.run([y1, y2])) np.testing.assert_allclose(*sess.run([log_det_y1, log_det_y2])) # inverse transform x1, log_det_x1 = flow1.inverse_transform(y1) x2, log_det_x2 = flow1.inverse_transform(y2) np.testing.assert_allclose(*sess.run([x1, x2])) np.testing.assert_allclose(*sess.run([log_det_x1, log_det_x2]))
def test_invert_flow(self): with self.test_session() as sess: # test invert a normal flow flow = QuadraticFlow(2., 5.) inv_flow = flow.invert() self.assertIsInstance(inv_flow, InvertFlow) self.assertEqual(inv_flow.x_value_ndims, 0) self.assertEqual(inv_flow.y_value_ndims, 0) self.assertFalse(inv_flow.require_batch_dims) test_x = np.arange(12, dtype=np.float32) + 1. test_y, test_log_det = quadratic_transform(npyops, test_x, 2., 5.) self.assertFalse(flow._has_built) y, log_det_y = inv_flow.inverse_transform(tf.constant(test_x)) self.assertTrue(flow._has_built) np.testing.assert_allclose(sess.run(y), test_y) np.testing.assert_allclose(sess.run(log_det_y), test_log_det) invertible_flow_standard_check(self, inv_flow, sess, test_y) # test invert an InvertFlow inv_inv_flow = inv_flow.invert() self.assertIs(inv_inv_flow, flow) # test use with FlowDistribution normal = Normal(mean=1., std=2.) inv_flow = QuadraticFlow(2., 5.).invert() distrib = FlowDistribution(normal, inv_flow) distrib_log_det = distrib.log_prob(test_x) np.testing.assert_allclose(*sess.run( [distrib_log_det, normal.log_prob(test_y) + test_log_det]))
def test_feature_shuffling_flow(self): np.random.seed(1234) with self.test_session() as sess: # axis = -1, value_ndims = 1 x = np.random.normal(size=[3, 4, 5, 6]).astype(np.float32) x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 6]) permutation = np.arange(6, dtype=np.int32) np.random.shuffle(permutation) y = x[..., permutation] log_det = np.zeros([3, 4, 5]).astype(np.float32) layer = FeatureShufflingFlow(axis=-1, value_ndims=1) y_out, log_det_out = layer.transform(x_ph) sess.run(tf.assign(layer._permutation, permutation)) y_out, log_det_out = sess.run( [y_out, log_det_out], feed_dict={x_ph: x}) np.testing.assert_equal(y_out, y) np.testing.assert_equal(log_det_out, log_det) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x}) assert_variables(['permutation'], trainable=False, scope='feature_shuffling_flow', collections=[tf.GraphKeys.MODEL_VARIABLES]) # axis = -2, value_ndims = 3 x = np.random.normal(size=[3, 4, 5, 6]).astype(np.float32) x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, 5, None]) permutation = np.arange(5, dtype=np.int32) np.random.shuffle(permutation) y = x[..., permutation, :] log_det = np.zeros([3]).astype(np.float32) layer = FeatureShufflingFlow(axis=-2, value_ndims=3) y_out, log_det_out = layer.transform(x_ph) sess.run(tf.assign(layer._permutation, permutation)) y_out, log_det_out = sess.run( [y_out, log_det_out], feed_dict={x_ph: x}) np.testing.assert_equal(y_out, y) np.testing.assert_equal(log_det_out, log_det) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x})
def test_invertible_dense(self): assert_allclose = functools.partial(np.testing.assert_allclose, rtol=1e-5) np.random.seed(1234) with self.test_session() as sess: x = np.random.normal(size=[3, 5, 7]).astype(np.float32) x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, 7]) kernel = np.random.normal(size=(7, 7)).astype(x.dtype) y, log_det = naive_invertible_linear(x, kernel, -1, 1) layer = InvertibleDense(strict_invertible=False) y_out, log_det_out = layer.transform(x_ph) # The kernel is initialized as an orthogonal matrix. As a result, # the log-det is zero (the computation result should be close to # zero, but may not be zero). This is not good for testing. # Thus we initialize it with an arbitrary matrix. sess.run(tf.assign(layer._kernel_matrix._matrix, kernel)) y_out, log_det_out = sess.run([y_out, log_det_out], feed_dict={x_ph: x}) assert_allclose(y_out, y) assert_allclose(log_det_out, log_det) invertible_flow_standard_check(self, layer, sess, x_ph, feed_dict={x_ph: x}, rtol=1e-5, atol=1e-6) assert_variables(['matrix'], trainable=True, scope='invertible_dense/kernel', collections=[tf.GraphKeys.MODEL_VARIABLES]) # test non-trainable with tf.Graph().as_default(): layer = InvertibleDense(strict_invertible=False, trainable=False) layer.apply(tf.zeros([2, 3, 4, 5])) assert_variables(['matrix'], trainable=False, scope='invertible_dense/kernel', collections=[tf.GraphKeys.MODEL_VARIABLES])
def check(x, x_value_ndims, y_value_shape, x_ph=None): flow = ReshapeFlow(x_value_ndims=x_value_ndims, y_value_shape=y_value_shape) if x_value_ndims > 0: batch_shape = list(x.shape[:-x_value_ndims]) else: batch_shape = list(x.shape) y_ans = np.reshape(x, batch_shape + list(y_value_shape)) log_det_ans = np.zeros(batch_shape) feed_dict = None if x_ph is not None: feed_dict = {x_ph: x} x = x_ph y, log_det = sess.run(flow.transform(x), feed_dict=feed_dict) np.testing.assert_allclose(y, y_ans) np.testing.assert_allclose(log_det, log_det_ans) invertible_flow_standard_check(self, flow, sess, x, feed_dict)
def check(x, y_shape, bs, channels_last=True, x_ph=None): # compute the answer y = naive_space_to_depth(x, bs, channels_last) log_det = np.zeros(y.shape[:-3], dtype=np.float32) self.assertTupleEqual(y.shape, y_shape) # get feed dict feed_dict = None if x_ph is not None: feed_dict = {x_ph: x} x = x_ph flow = SpaceToDepthFlow(bs, channels_last=channels_last) y_out, log_det_out = sess.run(flow.transform(x), feed_dict=feed_dict) np.testing.assert_allclose(y_out, y) np.testing.assert_allclose(log_det_out, log_det) invertible_flow_standard_check(self, flow, sess, x, feed_dict)
def test_different_value_ndims(self): def reshape_tail(x, value_ndims, shape): batch_shape = x.shape if value_ndims > 0: batch_shape = batch_shape[:-value_ndims] return np.reshape(x, batch_shape + tuple(shape)) def split_transform(x, split_axis, join_axis, x_value_ndims, y_shape, a1, b1, a2=None, b2=None): n1 = x.shape[split_axis] // 2 n2 = x.shape[split_axis] - n1 x1, x2 = np.split(x, [n1], axis=split_axis) y1, log_det1 = quadratic_transform(npyops, x1, a1, b1) if a2 is not None: y2, log_det2 = quadratic_transform(npyops, x2, a2, b2) else: y2, log_det2 = x2, np.zeros_like(x, dtype=np.float64) y1 = reshape_tail(y1, x_value_ndims, y_shape) y2 = reshape_tail(y2, x_value_ndims, y_shape) y = np.concatenate([y1, y2], axis=join_axis) if x_value_ndims > 0: reduce_axis = tuple(range(-x_value_ndims, 0)) log_det1 = np.sum(log_det1, axis=reduce_axis) log_det2 = np.sum(log_det2, axis=reduce_axis) log_det = log_det1 + log_det2 return y, log_det with self.test_session() as sess: np.random.seed(1234) x = 10. * np.random.normal(size=[3, 4, 5, 12]).astype(np.float64) # 2 -> 3, x_value_ndims = 3, y_value_ndims = 4 x_ph = tf.placeholder(dtype=tf.float64, shape=[None] * 4) flow = SplitFlow(split_axis=-2, join_axis=2, left=SequentialFlow([ QuadraticFlow(2., 5., value_ndims=3), ReshapeFlow(3, [4, -1, 2, 6]), ]), right=SequentialFlow([ QuadraticFlow(1.5, 3., value_ndims=3), ReshapeFlow(3, [4, -1, 2, 6]), ])) self.assertEqual(flow.x_value_ndims, 3) self.assertEqual(flow.y_value_ndims, 4) y, log_det = split_transform(x, split_axis=-2, join_axis=-3, x_value_ndims=3, y_shape=[4, -1, 2, 6], a1=2., b1=5., a2=1.5, b2=3.) y_out, log_det_out = sess.run(flow.transform(x_ph), feed_dict={x_ph: x}) np.testing.assert_allclose(y_out, y) np.testing.assert_allclose(log_det_out, log_det) invertible_flow_standard_check(self, flow, sess, x_ph, feed_dict={x_ph: x})
def test_equal_value_ndims(self): def split_transform(x, split_axis, value_ndims, a1, b1, a2=None, b2=None): n1 = x.shape[split_axis] // 2 n2 = x.shape[split_axis] - n1 x1, x2 = np.split(x, [n1], axis=split_axis) y1, log_det1 = quadratic_transform(npyops, x1, a1, b1) if a2 is not None: y2, log_det2 = quadratic_transform(npyops, x2, a2, b2) else: y2, log_det2 = x2, np.zeros_like(x, dtype=np.float64) y = np.concatenate([y1, y2], axis=split_axis) if value_ndims > 0: reduce_axis = tuple(range(-value_ndims, 0)) log_det1 = np.sum(log_det1, axis=reduce_axis) log_det2 = np.sum(log_det2, axis=reduce_axis) log_det = log_det1 + log_det2 return y, log_det with self.test_session() as sess: np.random.seed(1234) x = 10. * np.random.normal(size=[3, 4, 5, 6]).astype(np.float64) # static input, split_axis = -1, value_ndims = 1, right = None flow = SplitFlow(-1, QuadraticFlow(2., 5., value_ndims=1)) self.assertEqual(flow.x_value_ndims, 1) self.assertEqual(flow.y_value_ndims, 1) y, log_det = split_transform(x, split_axis=-1, value_ndims=1, a1=2., b1=5.) y_out, log_det_out = sess.run(flow.transform(x)) np.testing.assert_allclose(y_out, y) np.testing.assert_allclose(log_det_out, log_det) invertible_flow_standard_check(self, flow, sess, x, rtol=1e-4, atol=1e-5) # dynamic input, split_axis = -2, value_ndims = 2, right = None x_ph = tf.placeholder(dtype=tf.float64, shape=[None] * 4) flow = SplitFlow(-2, QuadraticFlow(2., 5., value_ndims=2)) self.assertEqual(flow.x_value_ndims, 2) self.assertEqual(flow.y_value_ndims, 2) y, log_det = split_transform(x, split_axis=-2, value_ndims=2, a1=2., b1=5.) y_out, log_det_out = sess.run(flow.transform(x_ph), feed_dict={x_ph: x}) np.testing.assert_allclose(y_out, y) np.testing.assert_allclose(log_det_out, log_det) invertible_flow_standard_check(self, flow, sess, x_ph, feed_dict={x_ph: x}) # dynamic input, split_axis = 2, value_ndims = 3 x_ph = tf.placeholder(dtype=tf.float64, shape=[None] * 4) flow = SplitFlow(split_axis=2, left=QuadraticFlow(2., 5., value_ndims=3), right=QuadraticFlow(1.5, 3., value_ndims=3)) self.assertEqual(flow.x_value_ndims, 3) self.assertEqual(flow.y_value_ndims, 3) y, log_det = split_transform(x, split_axis=2, value_ndims=3, a1=2., b1=5., a2=1.5, b2=3.) y_out, log_det_out = sess.run(flow.transform(x_ph), feed_dict={x_ph: x}) np.testing.assert_allclose(y_out, y) np.testing.assert_allclose(log_det_out, log_det) invertible_flow_standard_check(self, flow, sess, x_ph, feed_dict={x_ph: x})
def test_act_norm(self): assert_allclose = functools.partial( np.testing.assert_allclose, rtol=1e-5, atol=1e-6) np.random.seed(1234) x = np.random.normal(size=[3, 4, 5, 6, 7]) x_ph = tf.placeholder(dtype=tf.float64, shape=[None, None, 5, None, 7]) x2 = np.random.normal(size=[2, 3, 4, 5, 6, 7]) x2_ph = tf.placeholder(dtype=tf.float64, shape=[None, None, None, 5, None, 7]) x3 = np.random.normal(size=[4, 5, 6, 7]) x3_ph = tf.placeholder(dtype=tf.float64, shape=[None, 5, None, 7]) with self.test_session() as sess: # -- static input shape, scale_type = 'linear', value_ndims = 3 axis = [-1, -3] value_ndims = 3 var_shape = (5, 7) scale, bias, var_shape_aligned = naive_act_norm_initialize(x, axis) self.assertEqual(scale.shape, var_shape) self.assertEqual(bias.shape, var_shape) # test initialize act_norm = ActNorm(axis=axis, value_ndims=value_ndims, scale_type='linear') y_out, log_det_out = sess.run( act_norm.transform(tf.constant(x, dtype=tf.float64))) self.assertEqual(act_norm._bias.dtype.base_dtype, tf.float64) scale_out, bias_out = sess.run( [act_norm._pre_scale, act_norm._bias]) assert_allclose(scale_out, scale) assert_allclose(bias_out, bias) # test the transform output from the initializing procedure y, log_det = naive_act_norm_transform( x, var_shape_aligned, value_ndims, scale, bias) self.assertEqual(y.shape, x.shape) self.assertEqual(log_det.shape, x.shape[:-value_ndims]) assert_allclose(y_out, y) assert_allclose(log_det_out, log_det) # test use an initialized act_norm y2, log_det2 = naive_act_norm_transform( x2, var_shape_aligned, value_ndims, scale, bias) self.assertEqual(y2.shape, x2.shape) self.assertEqual(log_det2.shape, x2.shape[:-value_ndims]) y2_out, log_det2_out = sess.run( act_norm.transform(x2_ph), feed_dict={x2_ph: x2}) assert_allclose(y2_out, y2) assert_allclose(log_det2_out, log_det2) # invertible flow standard checks invertible_flow_standard_check(self, act_norm, sess, x) invertible_flow_standard_check( self, act_norm, sess, x2_ph, feed_dict={x2_ph: x2}) invertible_flow_standard_check( self, act_norm, sess, x3_ph, feed_dict={x3_ph: x3}) # -- dynamic input shape, scale_type = 'exp', value_ndims = 4 value_ndims = 4 # test initialize act_norm = ActNorm(axis=axis, value_ndims=value_ndims, scale_type='exp') y_out, log_det_out = sess.run( act_norm.transform(x_ph), feed_dict={x_ph: x}) self.assertEqual(act_norm._bias.dtype.base_dtype, tf.float64) scale_out, bias_out = sess.run( [tf.exp(act_norm._pre_scale), act_norm._bias]) assert_allclose(scale_out, scale) assert_allclose(bias_out, bias) # test the transform output from the initializing procedure y, log_det = naive_act_norm_transform( x, var_shape_aligned, value_ndims, scale, bias) self.assertEqual(y.shape, x.shape) self.assertEqual(log_det.shape, x.shape[:-value_ndims]) assert_allclose(y_out, y) assert_allclose(log_det_out, log_det) # test use an initialized act_norm y3, log_det3 = naive_act_norm_transform( x3, var_shape_aligned, value_ndims, scale, bias) self.assertEqual(y3.shape, x3.shape) self.assertEqual(log_det3.shape, x3.shape[:-value_ndims]) y3_out, log_det3_out = sess.run( act_norm.transform(x3_ph), feed_dict={x3_ph: x3}) assert_allclose(y3_out, y3) assert_allclose(log_det3_out, log_det3) # invertible flow standard checks invertible_flow_standard_check(self, act_norm, sess, x) invertible_flow_standard_check( self, act_norm, sess, x2_ph, feed_dict={x2_ph: x2}) invertible_flow_standard_check( self, act_norm, sess, x3_ph, feed_dict={x3_ph: x3})
def test_coupling_layer(self): assert_allclose = functools.partial( np.testing.assert_allclose, rtol=1e-5) np.random.seed(1234) kernel1 = np.random.normal(size=[3, 2]).astype(np.float32) kernel2 = np.random.normal(size=[2, 3]).astype(np.float32) shift1 = np.random.normal(size=[2]).astype(np.float32) shift2 = np.random.normal(size=[3]).astype(np.float32) def shift_and_scale_fn(x1, n2, no_scale=False): kernel = kernel1 if n2 == 2 else kernel2 shift = tf.convert_to_tensor(shift1 if n2 == 2 else shift2) assert(kernel.shape[-1] == n2) assert(shift.shape[-1] == n2) x1, s1, s2 = flatten_to_ndims(x1, 2) scale = unflatten_from_ndims(tf.matmul(x1, kernel), s1, s2) shift = shift + tf.zeros_like(scale, dtype=shift.dtype) if no_scale: scale = None return shift, scale def shift_and_scale_numpy_fn(x1, n2, no_scale=False): a, b = shift_and_scale_fn(x1, n2, no_scale=no_scale) if b is None: a = sess.run(a) else: a, b = sess.run([a, b]) return a, b with self.test_session() as sess: # test linear scale, primary x = np.random.normal(size=[3, 4, 5]).astype(np.float32) x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, 5]) axis = -1 value_ndims = 1 y_ans, log_det_ans = naive_coupling_layer( shift_and_scale_numpy_fn, x, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='linear', reverse=False ) layer = CouplingLayer( shift_and_scale_fn, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='linear' ) y, log_det = layer.transform(x_ph) y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x}) assert_allclose(y_out, y_ans) assert_allclose(log_det_out, log_det_ans) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x}) # test exp scale, primary axis = -1 value_ndims = 2 y_ans, log_det_ans = naive_coupling_layer( shift_and_scale_numpy_fn, x, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='exp', reverse=False ) layer = CouplingLayer( shift_and_scale_fn, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='exp' ) y, log_det = layer.transform(x_ph) y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x}) assert_allclose(y_out, y_ans) assert_allclose(log_det_out, log_det_ans) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x}) # test sigmoid scale, secondary sigmoid_scale_bias = np.exp(1) axis = -1 value_ndims = 1 y_ans, log_det_ans = naive_coupling_layer( shift_and_scale_numpy_fn, x, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='sigmoid', sigmoid_scale_bias=sigmoid_scale_bias, reverse=False ) layer = CouplingLayer( shift_and_scale_fn, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='sigmoid', sigmoid_scale_bias=sigmoid_scale_bias ) y, log_det = layer.transform(x_ph) y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x}) assert_allclose(y_out, y_ans) assert_allclose(log_det_out, log_det_ans) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x}) # test None scale, primary axis = -1 value_ndims = 1 y_ans, log_det_ans = naive_coupling_layer( functools.partial(shift_and_scale_numpy_fn, no_scale=True), x, axis=axis, value_ndims=value_ndims, secondary=False, scale_type=None, reverse=False ) layer = CouplingLayer( functools.partial(shift_and_scale_fn, no_scale=True), axis=axis, value_ndims=value_ndims, secondary=False, scale_type=None ) y, log_det = layer.transform(x_ph) y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x}) assert_allclose(y_out, y_ans) assert_allclose(log_det_out, log_det_ans) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x}) # test None scale, secondary axis = -1 value_ndims = 3 y_ans, log_det_ans = naive_coupling_layer( functools.partial(shift_and_scale_numpy_fn, no_scale=True), x, axis=axis, value_ndims=value_ndims, secondary=True, scale_type=None, reverse=False ) layer = CouplingLayer( functools.partial(shift_and_scale_fn, no_scale=True), axis=axis, value_ndims=value_ndims, secondary=True, scale_type=None ) y, log_det = layer.transform(x_ph) y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x}) assert_allclose(y_out, y_ans) assert_allclose(log_det_out, log_det_ans) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x})
def test_coupling_layer_with_conv2d(self): assert_allclose = functools.partial( np.testing.assert_allclose, atol=1e-5, rtol=5e-4) np.random.seed(1234) kernel1 = np.random.normal(size=[3, 3, 5, 6]).astype(np.float32) kernel2 = np.random.normal(size=[3, 3, 6, 5]).astype(np.float32) shift1 = np.random.normal(size=[6]).astype(np.float32) shift2 = np.random.normal(size=[5]).astype(np.float32) def shift_and_scale_fn(x1, n2, no_scale=False, channels_last=True): kernel = kernel1 if n2 == 6 else kernel2 shift = tf.convert_to_tensor(shift1 if n2 == 6 else shift2) assert (kernel.shape[-1] == n2) assert (shift.shape[-1] == n2) x1 = transpose_conv2d_channels_x_to_last( x1, channels_last=channels_last ) scale = conv2d(x1, n2, (3, 3), use_bias=False, kernel=kernel, channels_last=True) shift = shift + tf.zeros_like(scale, dtype=shift.dtype) scale = transpose_conv2d_channels_last_to_x(scale, channels_last) shift = transpose_conv2d_channels_last_to_x(shift, channels_last) if no_scale: scale = None return shift, scale def shift_and_scale_numpy_fn(x1, n2, no_scale=False, channels_last=True): a, b = shift_and_scale_fn(x1, n2, no_scale, channels_last) if b is None: a = sess.run(a) else: a, b = sess.run([a, b]) return a, b with self.test_session() as sess: # test exp scale, primary, NHWC x = np.random.normal(size=[11, 13, 32, 31, 11]).astype(np.float32) x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, None, None, 11]) axis = -1 value_ndims = 3 y_ans, log_det_ans = naive_coupling_layer( shift_and_scale_numpy_fn, x, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='exp', reverse=False ) layer = CouplingLayer( shift_and_scale_fn, axis=axis, value_ndims=value_ndims, secondary=False, scale_type='exp' ) y, log_det = layer.transform(x_ph) y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x}) assert_allclose(y_out, y_ans) assert_allclose(log_det_out, log_det_ans) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x}, rtol=5e-4, atol=1e-5 ) # test sigmoid scale, secondary, NCHW x = np.transpose(x, [0, 1, 4, 2, 3]) x_ph = tf.placeholder(dtype=tf.float32, shape=[None, None, 11, None, None]) axis = -3 value_ndims = 3 y_ans, log_det_ans = naive_coupling_layer( functools.partial(shift_and_scale_numpy_fn, channels_last=False), x, axis=axis, value_ndims=value_ndims, secondary=True, scale_type='sigmoid', reverse=False ) layer = CouplingLayer( functools.partial(shift_and_scale_fn, channels_last=False), axis=axis, value_ndims=value_ndims, secondary=True, scale_type='sigmoid' ) y, log_det = layer.transform(x_ph) y_out, log_det_out = sess.run([y, log_det], feed_dict={x_ph: x}) assert_allclose(y_out, y_ans) assert_allclose(log_det_out, log_det_ans) invertible_flow_standard_check( self, layer, sess, x_ph, feed_dict={x_ph: x}, rtol=5e-4, atol=1e-5 )