code: def wrap(state, ssize):
      return state.contiguous().view(ssize, -1, *state.shape[1:])
  
file_path: pytorch-struct/torch_struct/autoregressive.py
imports: import torch
  from .semirings import MaxSemiring, KMaxSemiring, TempMax
  from torch.distributions.distribution import Distribution
label_desc: Use image node_id 1 for calling a global function with example usage: wrap(state, ssize) and returns: state