code: def test_while_loop_nested():
      def _to_np_list(arrays):
          return [x.asnumpy() if x is not None else x for x in arrays]
  
      def _array(shape):
          return mx.nd.random.uniform(-1.0, 1.0, shape=shape)
  
      def inner_cond(i, j, x_sum, sc):
          return j < 2
  
      def inner_body(i, j, x_sum, sc):
          x_ij = sc.take(j).squeeze(axis=0)
          return (x_ij, x_ij), (i, j + 1, x_sum, sc)
  
      def outer_cond(i, j, x_sum, sc):
          return i < 2
  
      def outer_body(i, j, x_sum, sc):
          F = mx.sym if isinstance(i, mx.sym.Symbol) else mx.nd
          (x_ij, x_ji), (
              i_p,
              j_p,
              x_sum_p,
              sc_p,
          ) = F.contrib.while_loop(
              cond=inner_cond,
              func=inner_body,
              loop_vars=(i, j, x_sum, sc),
              max_iterations=2,
          )
          return (x_ij, x_ji), (i_p + 1, j_p - 2, x_sum_p, sc_p)
  
      def make_loop(i, j, x_sum, sc):
          F = mx.sym if isinstance(i, mx.sym.Symbol) else mx.nd
          (x_ij, x_ji), (
              new_i,
              new_j,
              x_sum_p,
              sc_p,
          ) = F.contrib.while_loop(
              cond=outer_cond,
              func=outer_body,
              loop_vars=(i, j, x_sum, sc),
              max_iterations=2,
          )
          return new_i, new_j, x_sum_p, sc_p, x_ij, x_ji
  
      args = {
          "i": mx.nd.array([0]),
          "j": mx.nd.array([0]),
          "x_sum": _array([5, 3]),
          "sc": _array([10, 10, 5, 3]),
      }
      args_grad = {
          "x_sum": _array([5, 3]),
          "sc": _array([10, 10, 5, 3]),
      }
      out_grad = [
          _array([1]),
          _array([1]),
          _array([5, 3]),
          _array([10, 10, 5, 3]),
          _array([2, 2, 10, 5, 3]),
          _array([2, 2, 10, 5, 3]),
      ]
  
      def _get_imp_result(is_train, args, args_grad, out_grad):
          args = {k: v.copy() for k, v in args.items()}
          args_grad = {k: v.copy() for k, v in args_grad.items()}
          i, j, x_sum, sc = [
              args[x].copy() for x in ["i", "j", "x_sum", "sc"]
          ]
          if is_train:
              x_sum.attach_grad()
              sc.attach_grad()
          with mx.autograd.record(train_mode=is_train):
              results = make_loop(i, j, x_sum, sc)
              cat_res = mx.nd.concat(
                  *[x.reshape(-1) for x in results], dim=0
              )
          if not is_train:
              return _to_np_list(results), []
          cat_grad = mx.nd.concat(
              *[x.reshape(-1) for x in out_grad], dim=0
          )
          assert cat_grad.shape == cat_res.shape
          cat_res.backward(out_grad=cat_grad)
          grads = [x_sum.grad, sc.grad]
          return _to_np_list(results), _to_np_list(grads)
  
      def _get_sym_result(is_train, args, args_grad, out_grad):
          args = {k: v.copy() for k, v in args.items()}
          args_grad = {k: v.copy() for k, v in args_grad.items()}
          i, j, x_sum, sc = [
              mx.sym.var("i"),
              mx.sym.var("j"),
              mx.sym.var("x_sum"),
              mx.sym.var("sc"),
          ]
          result_sym = mx.sym.Group(make_loop(i, j, x_sum, sc))
          executor = result_sym._bind(
              ctx=default_device(),
              args=args,
              args_grad=args_grad,
          )
          results = executor.forward(is_train=is_train)
          if not is_train:
              return _to_np_list(results), []
          executor.backward(out_grads=out_grad)
          grads = [
              executor.grad_dict["x_sum"],
              executor.grad_dict["sc"],
          ]
          return _to_np_list(results), _to_np_list(grads)
  
      for is_train in [True, False]:
          imp_out, imp_grad = _get_imp_result(
              is_train=is_train,
              args=args,
              args_grad=args_grad,
              out_grad=out_grad,
          )
          sym_out, sym_grad = _get_sym_result(
              is_train=is_train,
              args=args,
              args_grad=args_grad,
              out_grad=out_grad,
          )
          assert len(imp_out) == len(sym_out)
          assert len(imp_grad) == len(sym_grad)
          for x, y in zip(imp_out, sym_out):
              assert_almost_equal(x, y, rtol=1e-3, atol=1e-3)
          for x, y in zip(imp_grad, sym_grad):
              assert_almost_equal(x, y, 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 3 for calling a global function with example usage: test_while_loop_nested() and returns: unknown, mx, unknown, x_ij,  x_ij,  i,  unknown,  x_sum,  sc, unknown, x_ij,  x_ji,  unknown,  unknown,  x_sum_p,  sc_p, new_i,  new_j,  x_sum_p,  sc_p,  x_ij,  x_ji, _to_np_list,  _to_np_list, _to_np_list,  _to_np_list, _to_np_list,  list, _to_np_list,  list