code: def _tf_restore_batch_dims(x, num_nonbatch_dims, prototype):
    
    assert x.shape.ndims == 1 + num_nonbatch_dims
    new_shape = (
        prototype.shape.as_list()[:-num_nonbatch_dims] + x.shape.as_list()[1:])
    assert None not in new_shape
    if new_shape != x.shape.as_list():
      x = tf.reshape(x, new_shape)
    return x
docs: Reverse op of _tf_flatten_batch_dims.
  
    Un-flatten the first dimension of x to match all but the last
    num_nonbatch_dims dimensions of prototype.
  
    Args:
      x: a tf.Tensor with 1 + num_nonbatch_dims dimensions
      num_nonbatch_dims: an integer
      prototype: a tf.Tensor
  
    Returns:
      a tf.Tensor