code: def pad(nested_lists, pad_token, pad_shape=None):
      """Pad the input lists with the pad token. If pad_shape is provided, pad to that
      shape, otherwise infer the input shape and pad out to a square tensor shape.
      """
      if pad_shape is None:
          pad_shape = list(_infer_pad_shape(nested_lists))
      if not pad_shape:
          return nested_lists
      dimension, *rest = pad_shape
      result = [pad(nested, pad_token, rest) for nested in nested_lists]
      result += [_make_nested_padding(rest, pad_token)] * (
          dimension - len(result)
      )
      return result
  
imports: import itertools
  from collections import Counter
  from typing import Dict, List, Optional, Tuple
  import torch
  from pytext.common.constants import SpecialTokens, Token
  from pytext.utils import cuda, precision
label_desc: Use image node_id 4 for calling a global function with example usage: pad(nested_lists, pad_token, pad_shape) and returns: result, nested_lists