import tensorflow as tf

                    self.compute_shape(l2_shape[2], self.ff_pool_strides[1][1]),
                    self.compute_shape(l2_shape[3], self.ff_pool_strides[1][2]),
                    final_dim]
        else:
            l2_shape = tf.identity(x_shape)

        # Initialize hidden layer activities
        if self.hidden_init == 'identity':
            l1_h2 = tf.identity(x)
            l2_h2 = tf.zeros(l2_shape, dtype=self.dtype)
            l3_h2 = tf.zeros(l3_shape, dtype=self.dtype)
        elif self.hidden_init == 'random':
            l1_h2 = tf.random_normal(x_shape, dtype=self.dtype)
            l2_h2 = tf.random_normal(l2_shape, dtype=self.dtype)
            l3_h2 = tf.random_normal(l3_shape, dtype=self.dtype)
        elif self.hidden_init == 'zeros':
            l1_h2 = tf.zeros(x_shape, dtype=self.dtype)
            l2_h2 = tf.zeros(l2_shape, dtype=self.dtype)
            l3_h2 = tf.zeros(l3_shape, dtype=self.dtype)
        else:
            raise RuntimeError

        # While loop
        elems = [
            i0,
            x,
