import tensorflow as tf

  def loop_analysis(element):
    x = tf.expand_dims(element[0], 0)
    x_coori = tf.expand_dims(element[1], 0)
    y = analysis_transform(x_coori,x)
    return tf.squeeze(y,axis=0)

  element = [x,x_coori]
  ys = tf.map_fn(loop_analysis, element, dtype=tf.float32, parallel_iterations=1, back_prop=False)
  print("Analysis Transform")

  def loop_hyper_encoder(y):
    y = tf.expand_dims(y, 0)
    z = hyper_encoder(y)
    return tf.squeeze(z,axis=0)

  zs = tf.map_fn(loop_hyper_encoder, ys, dtype=tf.float32, parallel_iterations=1, back_prop=False)
  print("Hyper Encoder")

  z_hats, _ = entropy_bottleneck(zs, False)
  print("Quantize hyperprior")

  def loop_hyper_deocder(z):
    z = tf.expand_dims(z, 0)
    loc, scale = hyper_decoder(z)
    return tf.squeeze(loc, [0]), tf.squeeze(scale, [0])

