code: def standardize_shapes(features, batch_size=None):
    
    for fname in ["inputs", "targets"]:
      if fname not in features:
        continue
      f = features[fname]
      while len(f.get_shape()) < 4:
        f = tf.expand_dims(f, axis=-1)
      features[fname] = f
  
    if batch_size:
      # Ensure batch size is set on all features
      for _, t in six.iteritems(features):
        shape = t.get_shape().as_list()
        shape[0] = batch_size
        t.set_shape(t.get_shape().merge_with(shape))
        # Assert shapes are fully known
        t.get_shape().assert_is_fully_defined()
  
    return features
docs: Set the right shapes for the features.