code: def _find_domain_stack(op, parts):
      shape = broadcast_shape(*(x.shape for x in parts))
      dim = op.defaults["dim"]
      if dim >= 0:
          dim = dim - len(shape) - 1
      assert dim < 0
      split = dim + len(shape) + 1
      shape = shape[:split] + (len(parts),) + shape[split:]
      output = Array[parts[0].dtype, shape]
      return output
  
imports: import copyreg
  import functools
  import inspect
  import operator
  import warnings
  from functools import reduce
  from weakref import WeakValueDictionary
  import funsor.ops
  from funsor.ops.builtin import parse_ellipsis, parse_slice
  from funsor.util import broadcast_shape, get_backend, get_tracing_state, quote
label_desc: Use image node_id 20 for calling a global function with example usage: _find_domain_stack(op, parts) and returns: output