code: def map_fn_switch(fn, elems, use_map_fn=True, **kwargs):
    
    if use_map_fn:
      return tf.map_fn(fn, elems, **kwargs)
    elems_unpacked = (tf.unstack(e) for e in elems)
    out_unpacked = [fn(e) for e in zip(*elems_unpacked)]
    out = tf.stack(out_unpacked)
    return out
docs: Construct the graph with either tf.map_fn or a python for loop.
  
    This function is mainly for for benchmarking purpose.
  
    tf.map_fn is dynamic but is much slower than creating a static graph with
    for loop. However, having a for loop make the graph much longer to build
    and can consume too much RAM on distributed setting.
  
    Args:
      fn (fct): same that tf.map_fn but for now can only return a single tensor
        value (instead of a tuple of tensor for the general case)
      elems (tuple): same that tf.map_fn
      use_map_fn (bool): If True, tf.map_fn is used, if False, for _ in _: is used
        instead
      **kwargs: Additional tf.map_fn arguments (ignored if use_map_fn is False)
  
    Returns:
      tf.Tensor: the output of tf.map_fn