def shapes_match(a, b):
  """Recursively check if shapes of object `a` and `b` match.

  Will walk lists, tuples and dicts.

  Args:
    a: object of type (numpy.ndarray,tf.Tensor,list,tuple,dict)
        to check for matching shapes against `b`.
    b: object to check for matching shape against `a`.

  Returns:
    A boolean indicating whether the shapes of `a` and `b` match.
  """
  if isinstance(a, (tuple, list)) and isinstance(b, (tuple, list)):
    if len(a) != len(b):
      return False
    return all([shapes_match(ia, ib) for ia, ib in zip(a, b)])
  elif isinstance(a, dict) and isinstance(b, dict):
    if len(a) != len(b):
      return False
    match = True
    for (ak, av), (bk, bv) in zip(a.items(), b.items()):
      match = match and all([ak == bk and shapes_match(av, bv)])
    return match
  else:
    shape_checker = shape_checkers[(type(a), type(b))]
    return shape_checker(a, b)
Recursively check if shapes of object `a` and `b` match.

  Will walk lists, tuples and dicts.

  Args:
    a: object of type (numpy.ndarray,tf.Tensor,list,tuple,dict)
        to check for matching shapes against `b`.
    b: object to check for matching shape against `a`.

  Returns:
    A boolean indicating whether the shapes of `a` and `b` match.