def test_beta_share(self): """Test that beta share works correctly""" batch_size = 50 length = 10 test_1 = np.random.rand(batch_size, length) test_2 = np.random.rand(batch_size, length) with self.session() as sess: test_1 = tf.convert_to_tensor(test_1, dtype=tf.float32) test_2 = tf.convert_to_tensor(test_2, dtype=tf.float32) out_tensor = BetaShare()(test_1, test_2) sess.run(tf.global_variables_initializer()) out_tensor.eval() assert test_1.shape == out_tensor.shape assert test_2.shape == out_tensor.shape