code: def foreach(body, data, init_states, name="foreach"):
      """Run a for loop with user-defined computation over NDArrays on dimension 0.
  
      This operator simulates a for loop and body has the computation for an iteration
      of the for loop. It runs the computation in body on each slice from the input
      NDArrays.
  
      body takes two arguments as input and outputs a tuple of two elements,
      as illustrated below::
  
          out, states = body(data1, states)
  
      data1 can be either an NDArray or a list of NDArrays. If data is an NDArray,
      data1 is an NDArray. Otherwise, data1 is a list of NDArrays and has the same
      size as data. states is a list of NDArrays and have the same size as init_states.
      Similarly, out can be either an NDArray or a list of NDArrays, which are concatenated
      as the first output of foreach; states from the last execution of body
      are the second output of foreach.
  
      The computation done by this operator is equivalent to the pseudo code below
      when the input data is NDArray::
  
          states = init_states
          outs = []
          for i in data.shape[0]:
              s = data[i]
              out, states = body(s, states)
              outs.append(out)
          outs = stack(*outs)
  
  
      Parameters
      ----------
      body : HybridBlock.
          Define computation in an iteration.
      data: an NDArray or a list of NDArrays.
          The input data.
      init_states: an NDArray or nested lists of NDArrays.
          The initial values of the loop states.
  
      Returns
      -------
      outputs: an NDArray or nested lists of NDArrays.
          The output data concatenated from the output of all iterations.
      states: an NDArray or nested lists of NDArrays.
          The loop states in the last iteration.
  
      Examples
      --------
      >>> step = lambda data, states: (data + states[0], [states[0] * 2])
      >>> data = mx.np.random.uniform(size=(2, 10))
      >>> states = [mx.np.random.uniform(size=(10))]
      >>> outs, states = npx.control_flow.foreach(step, data, states)
      """
  
      def check_input(inputs, in_type, msg):
          is_NDArray_or_list = True
          if isinstance(inputs, list):
              for i in inputs:
                  if not isinstance(i, in_type):
                      is_NDArray_or_list = False
                      break
          else:
              is_NDArray_or_list = isinstance(inputs, in_type)
          assert is_NDArray_or_list, msg
  
      flatten_data, data_fmt = _flatten(data, "foreach input")
      check_input(
          flatten_data,
          np_ndarray,
          "data should be an mxnet.numpy.ndarray or a nested list of mxnet.numpy.ndarray",
      )
      flatten_state, state_fmt = _flatten(init_states, "foreach states")
      check_input(
          flatten_state,
          np_ndarray,
          "init_states should be an mxnet.numpy.ndarray or a nested list of mxnet.numpy.ndarray",
      )
  
      real_data = [
          ele[0].copy().detach() if ele is not None else None
          for ele in flatten_data
      ]
      real_state = [
          ele.copy().detach() if ele is not None else None
          for ele in flatten_state
      ]
  
      # If the input python function references to the symbols outside
      # the python function, we need to prune the computation graph constructed from
      # the function. One way of doing it is to mark the nodes in the computation graph
      # with AttrScope and prune the nodes without the special attribute.
      name = _get_unique_subgraph_name(name)
      with AttrScope(__subgraph_name__=name):
          data_names = [
              "data_subgraph{}".format(i)
              for i, ele in enumerate(real_data)
          ]
          state_names = [
              "state_subgraph{}".format(i)
              for i, ele in enumerate(real_state)
          ]
          symbol_data = [
              symbol.var(name).as_np_ndarray()
              for arg, name in zip(real_data, data_names)
          ]
          symbol_state = [
              symbol.var(name).as_np_ndarray()
              for arg, name in zip(real_state, state_names)
          ]
          dc.set_variable(real_data, symbol_data)
          dc.set_variable(real_state, symbol_state)
          in_eles = _regroup(real_data, data_fmt)
          in_states = _regroup(real_state, state_fmt)
          if dc.is_deferred_compute():
              out, states = body(in_eles, in_states)
          else:
              with ag.pause(), dc.context():
                  out, states = body(in_eles, in_states)
  
          flatten_out, out_fmt = _flatten(out, "foreach output")
          flatten_out_state, state_fmt = _flatten(
              states, "foreach loop_vars"
          )
  
          num_out_data = len(flatten_out)
          num_states = len(flatten_out_state)
          num_outputs = num_out_data + num_states
          sym_out = [
              dc.get_symbol(out_data) for out_data in flatten_out
          ]
          sym_states = [
              dc.get_symbol(out_state)
              for out_state in flatten_out_state
          ]
          dc.clear(flatten_out)
          dc.clear(flatten_out_state)
          g = _construct_subgraph(sym_out, sym_states)
  
      params_names = []
      params_data = []
      if hasattr(body, "collect_params"):
          for p in body.collect_params().values():
              params_names.append(p.var().name)
              params_data.append(p.data())
  
      subg_input_names = g.list_inputs()
  
      in_data, in_states, params = [], [], []
      in_data_locs, in_state_locs, remain_locs, in_state_index = (
          [],
          [],
          [],
          [],
      )
      for i, sub_name in enumerate(subg_input_names):
          if sub_name in data_names:
              in_data_locs.append(i)
              idx = data_names.index(sub_name)
              in_data.append(flatten_data[idx])
          elif sub_name in state_names:
              in_state_locs.append(i)
              idx = state_names.index(sub_name)
              in_states.append(flatten_state[idx])
              in_state_index.append(idx)
          elif sub_name in params_names:
              remain_locs.append(i)
              idx = params_names.index(sub_name)
              params.append(params_data[idx])
          else:
              raise AssertionError(
                  "the data arrays have to be used in the loop body"
              )
  
      ordered_ins = in_data + in_states + params
  
      ndoutput = _api_internal.foreach(
          g.handle,
          *ordered_ins,
          num_outputs,
          num_out_data,
          in_state_locs,
          in_data_locs,
          remain_locs,
          in_state_index,
      )
      if isinstance(ndoutput, NDArrayBase):
          ret = ndoutput
      else:
          ret = list(ndoutput)
      outs = []
      for i in range(num_outputs - num_states):
          outs.append(ret[i])
      outs = _regroup(outs, out_fmt)
      states = []
      for i in range(num_states):
          states.append(ret[num_outputs - num_states + i])
      states = _regroup(states, state_fmt)
  
      return (outs, states)
  
file_path: mxnet/python/mxnet/ndarray/numpy_extension/control_flow.py
imports: from .None import _api_internal
  from .._internal import NDArrayBase
  from ...util import set_module
  from ...numpy import ndarray
  from ...symbol import Symbol
  from ...base import _as_list
  from ...None import symbol, _deferred_compute, autograd
  from ...attribute import AttrScope, current
label_desc: Use image node_id 5 for calling a global function with example usage: foreach(body, data, init_states, name) and returns: outs,  states