code: def test_foreach_nested():
      # Test nested foreach.
      def step_in(in1, states):
          out = in1 * 2 + states[0]
          return (out, [out])
  
      def step_sym(in1, states):
          out1 = mx.sym.contrib.foreach(step_in, in1, states)
          out = mx.sym.broadcast_add(out1[0], states[0])
          return (
              out,
              [
                  mx.sym.squeeze(
                      mx.sym.slice(out, begin=(0, 0), end=(1, 2))
                  )
              ],
          )
  
      def step_nd(in1, states):
          out1 = mx.nd.contrib.foreach(step_in, in1, states)
          out = mx.nd.broadcast_add(out1[0], states[0])
          return (
              out,
              [
                  mx.nd.squeeze(
                      mx.nd.slice(out, begin=(0, 0), end=(1, 2))
                  )
              ],
          )
  
      data_sym = mx.sym.var("v1")
      state_sym = mx.sym.var("v2")
      out, states = mx.sym.contrib.foreach(
          step_sym, data_sym, [state_sym]
      )
      assert isinstance(states, list)
      assert len(states) == 1
      out = mx.sym.broadcast_add(out, states[0])
  
      js_1 = out.tojson()
      out = mx.sym.fromjson(js_1)
      js_2 = out.tojson()
      assert js_1 == js_2
  
      data = mx.nd.arange(8).reshape((2, 2, 2))
      state = mx.nd.arange(2)
      data_grad = mx.nd.empty(data.shape)
      state_grad = mx.nd.empty(state.shape)
      e = out._bind(
          ctx=default_device(),
          args={"v1": data, "v2": state},
          args_grad={"v1": data_grad, "v2": state_grad},
      )
      e.forward(is_train=True)
      out_grads = []
      for out in e.outputs:
          out_grads.append(mx.nd.random.uniform(shape=out.shape))
      e.backward(out_grads)
  
      data.attach_grad()
      state.attach_grad()
      with mx.autograd.record():
          out, states = mx.nd.contrib.foreach(step_nd, data, [state])
          assert isinstance(states, list)
          assert len(states) == 1
          res = mx.nd.broadcast_add(out, states[0])
      assert_almost_equal(
          res.asnumpy(), e.outputs[0].asnumpy(), rtol=1e-3, atol=1e-3
      )
  
      res.backward(out_grads[0])
      assert_almost_equal(
          data.grad.asnumpy(), data_grad.asnumpy(), rtol=1e-3, atol=1e-3
      )
      assert_almost_equal(
          state.grad.asnumpy(),
          state_grad.asnumpy(),
          rtol=1e-3,
          atol=1e-3,
      )
  
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 7 for calling a global function with example usage: test_foreach_nested() and returns: out,  list, out,  list, out,  list