code: def get_flatten_greedy_dims(tensor: torch.Tensor, max_ndim: int = 2):
      """get dims to flatten tensor up to max_ndim dimensions by merging small axes together"""
      dims = list(tensor.shape)
      while len(dims) > max_ndim:
          squeeze_ix = min(
              range(len(dims) - 1), key=lambda i: dims[i] * dims[i + 1]
          )
          squeezed_dim = dims.pop(squeeze_ix)
          dims[squeeze_ix] *= squeezed_dim
      return dims
  
label_desc: Use image node_id 2 for calling a global function with example usage: get_flatten_greedy_dims(tensor, max_ndim) and returns: dims