code: def test_foreach():
      v3 = mx.sym.var("v0")
      v4 = mx.sym.var("v1")
      v5 = mx.sym.var("v2")
      v6 = mx.sym.var("v3")
      v7 = mx.sym.var("v4")
      v8 = mx.sym.var("v5")
  
      def verify_foreach(
          step,
          in_syms,
          state_syms,
          free_syms,
          in_arrs,
          init_states,
          frees,
          out_grads,
          is_train=True,
          free_vars_func=None,
          num_iters=1,
      ):
          step_sym = lambda in_syms, state_syms: step(
              in_syms, state_syms, free_syms
          )
          res, states = mx.sym.contrib.foreach(
              step_sym, in_syms, state_syms
          )
          out = _as_list(res)
          num_outputs = len(out)
          for i in range(num_outputs):
              out[i] = out[i] * 2
          out.extend(states)
          out = mx.sym.Group(out)
          js_1 = out.tojson()
          out = mx.sym.fromjson(js_1)
          js_2 = out.tojson()
          assert js_1 == js_2
          arr_grads = []
          arg_dict = {}
          arg_grad_dict = {}
          i = 0
          for arr in _as_list(in_arrs):
              arr_grad = mx.nd.empty(arr.shape)
              arr_grads.append(arr_grad)
              arg_dict["v" + str(i)] = arr
              arg_grad_dict["v" + str(i)] = arr_grad
              i = i + 1
          for arr in init_states:
              arr_grad = mx.nd.empty(arr.shape)
              arr_grads.append(arr_grad)
              arg_dict["v" + str(i)] = arr
              arg_grad_dict["v" + str(i)] = arr_grad
              i = i + 1
          for arr in frees:
              arr_grad = mx.nd.empty(arr.shape)
              arr_grads.append(arr_grad)
              arg_dict["v" + str(i)] = arr
              arg_grad_dict["v" + str(i)] = arr_grad
              i = i + 1
  
          if is_train:
              e = out._bind(
                  ctx=default_device(),
                  args=arg_dict,
                  args_grad=arg_grad_dict,
              )
          else:
              e = out._bind(ctx=default_device(), args=arg_dict)
          # the inputs to forward and backward are the same so forward and backward
          # should always return the same outputs.
          for _ in range(num_iters):
              e.forward(is_train=is_train)
              if is_train:
                  # backward
                  tmp_grads = out_grads[0][:]
                  tmp_grads.extend(out_grads[1])
                  e.backward(tmp_grads)
  
          # Below we use imperative to reimplement foreach and compute its gradients.
          res = []
          for _ in range(len(_as_list(out_grads[0]))):
              res.append([])
          for arr in _as_list(in_arrs):
              arr.attach_grad()
          for arr in init_states:
              arr.attach_grad()
          for arr in frees:
              arr.attach_grad()
          with mx.autograd.record():
              frees_imp = (
                  frees
                  if free_vars_func is None
                  else free_vars_func(frees)
              )
              step_imp = lambda in_arrs, state_arrs: step(
                  in_arrs, state_arrs, frees_imp
              )
              states = [mx.nd.expand_dims(s, 0) for s in init_states]
              res, states = mx.nd.contrib.foreach(
                  step_imp, in_arrs, init_states
              )
  
              res2 = _as_list(res)
              for i in range(len(res2)):
                  res2[i] = res2[i] * 2
              outs = []
              outs[:] = res2[:]
              if isinstance(states, list):
                  outs.extend(states)
                  states = [mx.nd.expand_dims(s, 0) for s in states]
                  res2.extend(states)
              else:
                  outs.append(states)
                  states = mx.nd.expand_dims(states, 0)
                  res2.append(states)
              if is_train:
                  res = mx.nd.concat(*res2, dim=0)
  
          tmp_grads = out_grads[0][:]
          tmp_grads1 = [
              mx.nd.expand_dims(grad, 0) for grad in out_grads[1]
          ]
          tmp_grads.extend(tmp_grads1)
          if is_train:
              res.backward(mx.nd.concat(*tmp_grads, dim=0))
          for i in range(len(outs)):
              assert e.outputs[i].shape == outs[i].shape
              assert_almost_equal(
                  e.outputs[i].asnumpy(),
                  outs[i].asnumpy(),
                  rtol=1e-3,
                  atol=1e-3,
              )
          if is_train:
              all_ins = _as_list(in_arrs)[:]
              all_ins.extend(init_states)
              all_ins.extend(frees)
              size = min(len(all_ins), len(e.grad_arrays))
              for i in range(size):
                  assert_almost_equal(
                      all_ins[i].grad.asnumpy(),
                      e.grad_arrays[i].asnumpy(),
                      rtol=1e-3,
                      atol=1e-3,
                  )
  
      # Test cases:
      # * graph inputs are stored in different orders.
      #   This is to test if foreach finds the data arrays and weight arrays
      #   in the right location.
      # * the number of iterations: odd or even.
      # * multiple inputs and multiple outputs.
      # * inference.
      def step1(in1, states, free):
          out = in1 * 2 + states[0] + free[0]
          return (out, [out])
  
      frees1 = [mx.nd.arange(2), mx.nd.arange(2) + 1]
      arrs = mx.nd.arange(6).reshape(shape=(3, 2))
      states = [mx.nd.arange(2)]
      out_grads = [
          [mx.nd.random.uniform(-10, 10, arrs.shape)],
          [mx.nd.random.uniform(-10, 10, states[0].shape)],
      ]
      verify_foreach(
          step1,
          v3,
          [v4],
          [v5 + v6],
          arrs,
          states,
          frees1,
          out_grads,
          True,
          lambda frees: [frees[0] + frees[1]],
      )
      verify_foreach(
          step1,
          v3,
          [v4],
          [v5 + v6],
          arrs,
          states,
          frees1,
          out_grads,
          False,
          lambda frees: [frees[0] + frees[1]],
      )
      verify_foreach(
          step1,
          v3,
          [v4],
          [v5 + v6],
          arrs,
          states,
          frees1,
          out_grads,
          True,
          lambda frees: [frees[0] + frees[1]],
          5,
      )
      verify_foreach(
          step1,
          v3,
          [v4],
          [v5 + v6],
          arrs,
          states,
          frees1,
          out_grads,
          False,
          lambda frees: [frees[0] + frees[1]],
          5,
      )
  
      # Test the even number of iterations.
      frees = [mx.nd.random.uniform(shape=(2))]
      arrs = mx.nd.random.uniform(shape=(2, 2))
      out_grads = [
          [mx.nd.random.uniform(-10, 10, arrs.shape)],
          [mx.nd.random.uniform(-10, 10, states[0].shape)],
      ]
      verify_foreach(
          step1, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step1, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
      # Test the odd number of iterations
      arrs = mx.nd.random.uniform(shape=(3, 2))
      out_grads = [
          [mx.nd.random.uniform(-10, 10, arrs.shape)],
          [mx.nd.random.uniform(-10, 10, states[0].shape)],
      ]
      verify_foreach(
          step1, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step1, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
  
      # Reorder the input and state in the subgraph inputs.
      def step2(in1, states, free):
          out = states[0] + in1 * 2 + free[0]
          return (out, [out])
  
      # Test the even number of iterations.
      arrs = mx.nd.random.uniform(shape=(2, 2))
      out_grads = [
          [mx.nd.random.uniform(-10, 10, arrs.shape)],
          [mx.nd.random.uniform(-10, 10, states[0].shape)],
      ]
      verify_foreach(
          step2, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step2, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
      # Test the odd number of iterations.
      arrs = mx.nd.random.uniform(shape=(3, 2))
      out_grads = [
          [mx.nd.random.uniform(-10, 10, arrs.shape)],
          [mx.nd.random.uniform(-10, 10, states[0].shape)],
      ]
      verify_foreach(
          step2, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step2, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
  
      # Test multiple inputs and outputs.
      def step3(in1, states, free):
          out = (
              in1[0] + in1[1] * 2 + states[0] + states[1] * 2 + free[0]
          )
          return ([out, out], [out * 2, out * 3])
  
      arrs = [
          mx.nd.random.uniform(shape=(3, 2)),
          mx.nd.random.uniform(shape=(3, 2)),
      ]
      states = [
          mx.nd.random.uniform(shape=(2)),
          mx.nd.random.uniform(shape=(2)),
      ]
      out_grads = [
          [
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
              mx.nd.random.uniform(-10, 10, arrs[1].shape),
          ],
          [
              mx.nd.random.uniform(-10, 10, states[0].shape),
              mx.nd.random.uniform(-10, 10, states[1].shape),
          ],
      ]
      verify_foreach(
          step3,
          [v3, v4],
          [v5, v6],
          [v7],
          arrs,
          states,
          frees,
          out_grads,
      )
      verify_foreach(
          step3,
          [v3, v4],
          [v5, v6],
          [v7],
          arrs,
          states,
          frees,
          out_grads,
          False,
      )
  
      # Test multiple inputs and outputs.
      # The order of subgraph inputs doesn't match the operator inputs
      def step4(in1, states, free):
          out = (
              in1[1] * 2 + states[0] + free[0] + states[1] * 2 + in1[0]
          )
          return ([out, out * 2], [out * 2, out * 3])
  
      arrs = [
          mx.nd.random.uniform(shape=(3, 2)),
          mx.nd.random.uniform(shape=(3, 2)),
      ]
      states = [
          mx.nd.random.uniform(shape=(2)),
          mx.nd.random.uniform(shape=(2)),
      ]
      out_grads = [
          [
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
              mx.nd.random.uniform(-10, 10, arrs[1].shape),
          ],
          [
              mx.nd.random.uniform(-10, 10, states[0].shape),
              mx.nd.random.uniform(-10, 10, states[1].shape),
          ],
      ]
      verify_foreach(
          step4,
          [v3, v4],
          [v5, v6],
          [v7],
          arrs,
          states,
          frees,
          out_grads,
      )
      verify_foreach(
          step4,
          [v3, v4],
          [v5, v6],
          [v7],
          arrs,
          states,
          frees,
          out_grads,
          False,
      )
  
      # Test multiple inputs and outputs.
      # The data inputs and states have different shapes.
      def step5(in1, states, free):
          if isinstance(in1[0], mx.nd.NDArray):
              out1 = mx.nd.broadcast_add(
                  states[0] + free[1], in1[1] * 2
              )
              out2 = mx.nd.broadcast_add(
                  in1[0], free[0] + states[1] * 2
              )
          else:
              out1 = mx.sym.broadcast_add(
                  states[0] + free[1], in1[1] * 2
              )
              out2 = mx.sym.broadcast_add(
                  in1[0], free[0] + states[1] * 2
              )
          return ([out1, out2 * 2], [states[0] * 2, states[1] * 3])
  
      frees = [
          mx.nd.random.uniform(shape=(2)),
          mx.nd.random.uniform(shape=(2, 2)),
      ]
      arrs = [
          mx.nd.random.uniform(shape=(3, 2, 2)),
          mx.nd.random.uniform(shape=(3, 2)),
      ]
      states = [
          mx.nd.random.uniform(shape=(2, 2)),
          mx.nd.random.uniform(shape=(2)),
      ]
      out_grads = [
          [
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
          ],
          [
              mx.nd.random.uniform(-10, 10, states[0].shape),
              mx.nd.random.uniform(-10, 10, states[1].shape),
          ],
      ]
      verify_foreach(
          step5,
          [v3, v4],
          [v5, v6],
          [v7, v8],
          arrs,
          states,
          frees,
          out_grads,
          False,
      )
  
      # Test multiple inputs and outputs.
      # The data inputs and states have different shapes and data types.
      def step6(in1, states, free):
          if isinstance(in1[0], mx.nd.NDArray):
              out1 = mx.nd.broadcast_add(
                  states[0] + mx.nd.cast(free[1], "float32"),
                  mx.nd.cast(in1[1], "float32") * 2,
              )
              out2 = mx.nd.broadcast_add(
                  in1[0], free[0] + mx.nd.cast(states[1], "float32") * 2
              )
          else:
              out1 = mx.sym.broadcast_add(
                  states[0] + mx.sym.cast(free[1], "float32"),
                  mx.sym.cast(in1[1], "float32") * 2,
              )
              out2 = mx.sym.broadcast_add(
                  in1[0],
                  free[0] + mx.sym.cast(states[1], "float32") * 2,
              )
          return ([out1, out2 * 2], [states[0] * 2, states[1] * 3])
  
      frees = [
          mx.nd.random.uniform(shape=(2)),
          mx.nd.cast(mx.nd.random.uniform(shape=(2, 2)), "float64"),
      ]
      arrs = [
          mx.nd.random.uniform(shape=(3, 2, 2)),
          mx.nd.cast(
              mx.nd.random.uniform(shape=(3, 2)), dtype="float16"
          ),
      ]
      states = [
          mx.nd.random.uniform(shape=(2, 2)),
          mx.nd.cast(mx.nd.random.uniform(shape=(2)), dtype="int32"),
      ]
      out_grads = [
          [
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
          ],
          [
              mx.nd.random.uniform(-10, 10, states[0].shape),
              mx.nd.random.uniform(-10, 10, states[1].shape),
          ],
      ]
      verify_foreach(
          step6,
          [v3, v4],
          [v5, v6],
          [v7, v8],
          arrs,
          states,
          frees,
          out_grads,
          False,
      )
  
      # Test multiple inputs and outputs.
      # some of the inputs are used twice.
      def step7(in1, states, free):
          out1 = states[0] + in1[0] + free[1] + in1[1] * 2 + free[0]
          out2 = in1[0] + free[0] + states[1] * 2 + in1[1]
          return ([out1, out2 * 2], [states[0] * 2, states[1] * 3])
  
      frees = [
          mx.nd.random.uniform(shape=(2)),
          mx.nd.random.uniform(shape=(2)),
      ]
      arrs = [
          mx.nd.random.uniform(shape=(3, 2)),
          mx.nd.random.uniform(shape=(3, 2)),
      ]
      states = [
          mx.nd.random.uniform(shape=(2)),
          mx.nd.random.uniform(shape=(2)),
      ]
      out_grads = [
          [
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
              mx.nd.random.uniform(-10, 10, arrs[0].shape),
          ],
          [
              mx.nd.random.uniform(-10, 10, states[0].shape),
              mx.nd.random.uniform(-10, 10, states[1].shape),
          ],
      ]
      verify_foreach(
          step7,
          [v3, v4],
          [v5, v6],
          [v7, v8],
          arrs,
          states,
          frees,
          out_grads,
          False,
      )
  
      # Test the case that the output is the input.
      arrs = mx.nd.random.uniform(shape=(3, 2))
      states = [mx.nd.arange(2)]
      frees = [mx.nd.random.uniform(shape=(2))]
      out_grads = [
          [mx.nd.random.uniform(-10, 10, arrs.shape)],
          [mx.nd.random.uniform(-10, 10, states[0].shape)],
      ]
  
      def step8(in1, states, free):
          return (in1, [states[0] * free[0]])
  
      verify_foreach(
          step8, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step8, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
  
      def step9(in1, states, free):
          return (in1 * free[0], states)
  
      verify_foreach(
          step9, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step9, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
  
      # Test the case that not all inputs are used.
      def step10(in1, states, free):
          return (in1, states)
  
      verify_foreach(
          step10, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step10, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
  
      def step11(in1, states, free):
          return (in1, free)
  
      try:
          verify_foreach(
              step11, v3, [v4], [v5], arrs, states, frees, out_grads
          )
          verify_foreach(
              step11,
              v3,
              [v4],
              [v5],
              arrs,
              states,
              frees,
              out_grads,
              False,
          )
      except AssertionError:
          print("the states have to be used")
  
      def step12(in1, states, free):
          return (in1, [states[0] + 1, states[0] + 2])
  
      states = [
          mx.nd.random.uniform(shape=(2)),
          mx.nd.random.uniform(shape=(2)),
      ]
      frees = []
      try:
          verify_foreach(
              step12, v3, [v4, v5], [], arrs, states, frees, out_grads
          )
          verify_foreach(
              step12,
              v3,
              [v4, v5],
              [],
              arrs,
              states,
              frees,
              out_grads,
              False,
          )
      except AssertionError:
          print("the states have to be used")
  
      # test without free variables.
      def step13(in1, states, free):
          return (in1, states)
  
      states = [mx.nd.random.uniform(shape=(2))]
      verify_foreach(step13, v3, [v4], [], arrs, states, [], out_grads)
      verify_foreach(
          step13, v3, [v4], [], arrs, states, [], out_grads, False
      )
  
      # test when there isn't output data or output states.
      def step14(in1, states, free):
          return (in1 + free[0], [])
  
      frees = [mx.nd.random.uniform(shape=(2))]
      out_grads = [[mx.nd.random.uniform(-10, 10, arrs.shape)], []]
      verify_foreach(step14, v3, [], [v4], arrs, [], frees, out_grads)
      verify_foreach(
          step14, v3, [], [v4], arrs, [], frees, out_grads, False
      )
  
      def step15(in1, states, free):
          return ([], [in1 * states[0] * free[0]])
  
      out_grads = [[], [mx.nd.random.uniform(-10, 10, states[0].shape)]]
      verify_foreach(
          step15, v3, [v4], [v5], arrs, states, frees, out_grads
      )
      verify_foreach(
          step15, v3, [v4], [v5], arrs, states, frees, out_grads, False
      )
  
      # Test the case of iterating on a 1D data array.
      def step16(in1, states, free):
          return ([in1[0] * states[0]], [states[0] * 2])
  
      arrs = [mx.nd.arange(3)]
      states = [mx.nd.random.uniform(shape=(1))]
      out_grads = [
          [mx.nd.random.uniform(-10, 10, (3, 1))],
          [mx.nd.random.uniform(-10, 10, (1))],
      ]
      verify_foreach(
          step16, [v3], [v4], [], arrs, states, [], out_grads
      )
      verify_foreach(
          step16, [v3], [v4], [], arrs, states, [], out_grads, False
      )
  
      def step17(in1, states, free):
          return ([in1[1] * in1[0] * states[0]], [states[0] * 2])
  
      arrs = [mx.nd.random.uniform(shape=(3, 1)), mx.nd.arange(3)]
      states = [mx.nd.random.uniform(shape=(1))]
      out_grads = [
          [mx.nd.random.uniform(-10, 10, (3, 1))],
          [mx.nd.random.uniform(-10, 10, (1))],
      ]
      verify_foreach(
          step17, [v3, v4], [v5], [], arrs, states, [], out_grads
      )
      verify_foreach(
          step17, [v3, v4], [v5], [], arrs, states, [], out_grads, False
      )
  
file_path: mxnet/tests/python/unittest/test_contrib_control_flow.py
imports: import pytest
  import mxnet
  from numpy.testing import assert_allclose
  from mxnet.test_utils import *
  from mxnet.base import _as_list
label_desc: Use image node_id 6 for calling a global function with example usage: test_foreach() and returns: out,  list, out,  list, list,  list, list,  list, list,  list, list,  list, list,  list, in1,  list, unknown,  states, in1,  states, in1,  free, in1,  list, in1,  states, unknown,  list, list,  list, list,  list, list,  list