code: def test_while_loop_for_foreach():
      def make_true_cond():
          return lambda loop_vars, _: (loop_vars[0] < 1e35).prod()
  
      def make_false_cond():
          return lambda loop_vars, _: (loop_vars[0] > 1e35).prod()
  
      def make_for_cond(length):
          return lambda loop_vars, _: loop_vars[0] < length
  
      def case_0():
          # This is a simple testcase that all loop steps are independent'
          # It basically scans the array and outputs itself
          # There is 1 output
          # There is 1 state: i
          def _simple_func(loop, free):
              (i,), (scanned,) = loop, free
              in_ = scanned.take(i).squeeze(axis=0)
              return (in_, i + 1)
  
          _verify_while_loop(
              cond=make_true_cond(),
              func=_simple_func,
              max_iterations=1,
              is_train=True,
              is_for=True,
              loop_var_shapes=[
                  (1,),  # i
              ],
              free_var_shapes=[
                  (1, 3),  # scanned
              ],
              n_steps=1,
          )
  
      def case_1(**params):
          # This is a simple testcase that simulates a cumulative sum
          # There is 1 output
          # There is 1 state: s
          step_funcs = [
              lambda a, b, s: s,
              lambda a, b, s: a * 1.5 + b * 2.5 - s * 3.5,
              lambda a, b, s: a * 1.5 - s * 3.5 + b * 2.5,
              lambda a, b, s: b * 2.5 + a * 1.5 - s * 3.5,
              lambda a, b, s: b * 2.5 - s * 3.5 + a * 1.5,
              lambda a, b, s: s * -3.5 + a * 1.5 + b * 2.5,
              lambda a, b, s: s * -3.5 + b * 2.5 + a * 1.5,
              lambda a, b, s: a * 2.5 * b + s * 0.3,
              lambda a, b, s: b * 2.5 * a + s * 0.3,
              lambda a, b, s: 2.5 * a * b + s * 0.3,
              lambda a, b, s: b * a * 2.5 + s * 0.3,
              lambda a, b, s: 2.5 * b * a + s * 0.3,
              lambda a, b, s: b * a * 2.5 + s * 0.3,
              lambda a, b, s: s * 0.3 + a * 2.5 * b,
              lambda a, b, s: s * 0.3 + b * 2.5 * a,
              lambda a, b, s: s * 0.3 + 2.5 * a * b,
              lambda a, b, s: s * 0.3 + b * a * 2.5,
              lambda a, b, s: s * 0.3 + 2.5 * b * a,
              lambda a, b, s: s * 0.3 + b * a * 2.5,
          ]
  
          def make_func(step_func):
              def step(loop, free):
                  (s,), (a, b) = loop, free
                  out = step_func(a, b, s)
                  return (out, out)
  
              return step
  
          case_id = 0
          for is_train in [True, False]:
              for step_func in step_funcs:
                  case_id += 1
                  _verify_while_loop(
                      func=make_func(step_func),
                      is_train=is_train,
                      is_for=False,
                      **params,
                  )
  
      def case_2(**params):
          # This is a testcase that involves non-differentiable operators
          # There is 1 output
          # There is 2 states: i, s
          step_funcs = [
              lambda in_, s, f_1: (in_ * 2) * s * f_1,
              lambda in_, s, f_1: (in_ * 2) * f_1 * s,
              lambda in_, s, f_1: s * (in_ * 2) * f_1,
              lambda in_, s, f_1: s * f_1 * (in_ * 2),
              lambda in_, s, f_1: f_1 * (in_ * 2) * s,
              lambda in_, s, f_1: f_1 * s * (in_ * 2),
              lambda in_, s, f_1: (2 * in_) * s * f_1,
              lambda in_, s, f_1: (2 * in_) * f_1 * s,
              lambda in_, s, f_1: s * (2 * in_) * f_1,
              lambda in_, s, f_1: s * f_1 * (2 * in_),
              lambda in_, s, f_1: f_1 * (2 * in_) * s,
              lambda in_, s, f_1: f_1 * s * (2 * in_),
          ]
  
          def make_func(step_func):
              """This simulates:
              def compute(s, inputs, f_1, length):
                  outputs = []
                  for i in range(length):
                      s += inputs[i] * 2 + f_1
                      outputs.append(s)
                  return outputs, s
              """
  
              def step(loop, free):
                  (i, s), (scanned, f_1, _) = loop, free
                  in_ = scanned.take(i).squeeze(axis=0)
                  out = step_func(in_, s, f_1)
                  return (out, (i + 1, out))
  
              return step
  
          case_id = 0
          for is_train in [True, False]:
              for step_func in step_funcs:
                  case_id += 1
                  _verify_while_loop(
                      func=make_func(step_func),
                      max_iterations=1000,
                      is_train=is_train,
                      is_for=True,
                      **params,
                  )
  
      def case_3(length, **params):
          # This is a testcase for multiple non-differentiable operators and different ways of slicing
          # There are 2 outputs
          # There are 3 states: i, s_0, s_1
          step_funcs = [
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_1,
              lambda i_0, i_1, s_0, s_1, f_0: s_0,
              lambda i_0, i_1, s_0, s_1, f_0: s_1,
              lambda i_0, i_1, s_0, s_1, f_0: f_0,
          ]
  
          def make_func(step_func):
              """This simulates:
              def compute(input_0, input_1, s_0, s_1, f_0, length):
                  output_0 = []
                  output_1 = []
                  for i in range(length):
                      i_0 = input_0[i]
                      i_1 = input_1[length - 1 - i]
                      out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0
                      s_0 = (s_0 + out) * 1.05
                      s_1 = (s_1 - out * 0.5) * 0.95
                      output_0.append(out)
                      output_1.append(out * 1.5)
                  return outputs, s_0, s_1
              """
  
              def step(loop, free):
                  (i, s_0, s_1), (sc_0, sc_1, f_0, _) = loop, free
                  i_0 = sc_0.take(i).squeeze(axis=0)
                  i_1 = sc_1.take(length - 1 - i).squeeze(axis=0)
                  out = step_func(i_0, i_1, s_0, s_1, f_0)
                  return (
                      [out, out * 1.5],
                      [
                          i + 1,
                          (s_0 + out) * 1.05,
                          (s_1 - out * 0.5) * 0.95,
                      ],
                  )
  
              return step
  
          case_id = 0
          for is_train in [True, False]:
              for step_func in step_funcs:
                  case_id += 1
                  _verify_while_loop(
                      func=make_func(step_func),
                      max_iterations=1000,
                      is_train=is_train,
                      is_for=True,
                      **params,
                  )
  
      def case_4(length, single_shape, **params):
          # It is for the case that inputs & outputs are the same
          # There are 3 outputs
          # There are 4 states: i, s_0, s_1, s_2
          # i is used in both non-differentiable (take) and differentiable (+) occasions
          step_funcs = [
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_1,
              lambda i_0, i_1, s_0, s_1, f_0: s_0,
              lambda i_0, i_1, s_0, s_1, f_0: s_1,
              lambda i_0, i_1, s_0, s_1, f_0: f_0,
          ]
  
          def make_func(step_func):
              """This simulates:
              def compute(input_0, input_1, s_0, s_1, s_2, f_0, length):
                  # here s_2 remains untouched
                  output_0 = []
                  output_1 = []
                  output_2 = []
                  for i in range(length):
                      i_0 = input_0[i]
                      i_1 = input_1[length - 1 - i]
                      out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0
                      out = out * i * i_0 * i_1
                      s_0 = (s_0 + out) * 1.05
                      s_1 = (s_1 - out * 0.5) * 0.95
                      output_0.append(out)
                      output_1.append(f_0)
                      output_2.append(out * 1.5)
                  return output_0, output_1, output_2, s_0, s_1, s_2
              """
  
              def step(loop, free):
                  (i, s_0, s_1, s_2), (sc_0, sc_1, f_0, _) = loop, free
                  i_0 = sc_0.take(i).squeeze(axis=0)
                  i_1 = sc_1.take(length - 1 - i).squeeze(axis=0)
                  out = step_func(i_0, i_1, s_0, s_1, f_0)
                  out = out * i.reshape(
                      [1] * len(single_shape)
                  ).broadcast_to(single_shape)
                  out = out * i_0 * i_1
                  return (
                      [out, f_0, out * 1.5],
                      [
                          i + 1,
                          (s_0 + out) * 1.05,
                          (s_1 - out * 0.5) * 0.95,
                          s_2,
                      ],
                  )
  
              return step
  
          case_id = 0
          for is_train in [True, False]:
              for step_func in step_funcs:
                  case_id += 1
                  _verify_while_loop(
                      func=make_func(step_func),
                      max_iterations=1000,
                      is_train=is_train,
                      is_for=True,
                      **params,
                  )
  
      def case_5(length, single_shape, **params):
          # It is for the case that inputs & outputs are the same
          # There are 0 outputs
          # There are 4 states: i, s_0, s_1, s_2
          # i is used in both differentiable (take) and non-differentiable (+) occasions
          step_funcs = [
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_1,
              lambda i_0, i_1, s_0, s_1, f_0: s_0,
              lambda i_0, i_1, s_0, s_1, f_0: s_1,
              lambda i_0, i_1, s_0, s_1, f_0: f_0,
          ]
  
          def make_func(step_func):
              """This simulates:
              def compute(input_0, input_1, s_0, s_1, s_2, f_0, length):
                  # here s_2 remains untouched
                  output_0 = []
                  output_1 = []
                  output_2 = []
                  for i in range(length):
                      i_0 = input_0[i]
                      i_1 = input_1[length - 1 - i]
                      out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0
                      out = out * i * i_0 * i_1
                      s_0 = (s_0 + out) * 1.05
                      s_1 = (s_1 - out * 0.5) * 0.95
                      output_0.append(out)
                      output_1.append(f_0)
                      output_2.append(out * 1.5)
                  return output_0, output_1, output_2, s_0, s_1, s_2
              """
  
              def step(loop, free):
                  (i, s_0, s_1, s_2), (sc_0, sc_1, f_0, _) = loop, free
                  i_0 = sc_0.take(i).squeeze(axis=0)
                  i_1 = sc_1.take(length - 1 - i).squeeze(axis=0)
                  out = step_func(i_0, i_1, s_0, s_1, f_0)
                  out = out * i.reshape(
                      [1] * len(single_shape)
                  ).broadcast_to(single_shape)
                  out = out * i_0 * i_1
                  return (
                      [],
                      [
                          i + 1,
                          (s_0 + out) * 1.05,
                          (s_1 - out * 0.5) * 0.95,
                          s_2,
                      ],
                  )
  
              return step
  
          case_id = 0
          for is_train in [True, False]:
              for step_func in step_funcs:
                  case_id += 1
                  _verify_while_loop(
                      func=make_func(step_func),
                      max_iterations=1000,
                      is_train=is_train,
                      is_for=True,
                      **params,
                  )
  
      def case_6(length, single_shape, **params):
          # It is for the case that inputs & outputs are the same
          # There are 3 outputs
          # There are 4 states: i, s_0, s_1, s_2
          # i is used in both differentiable (take) and non-differentiable (+) occasions
          step_funcs = [
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0
              * (i_1 * 2)
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * (s_1 * 2)
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * s_0
              * f_0
              * (s_1 * 2),
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * s_0
              * f_0,
              lambda i_0, i_1, s_0, s_1, f_0: (i_1 * 2)
              * i_0
              * (s_1 * 2)
              * f_0
              * s_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_0,
              lambda i_0, i_1, s_0, s_1, f_0: i_1,
              lambda i_0, i_1, s_0, s_1, f_0: s_0,
              lambda i_0, i_1, s_0, s_1, f_0: s_1,
              lambda i_0, i_1, s_0, s_1, f_0: f_0,
          ]
  
          def make_func(step_func):
              """This simulates:
              def compute(input_0, input_1, s_0, s_1, s_2, f_0, length):
                  # here s_2 remains untouched
                  output_0 = []
                  output_1 = []
                  output_2 = []
                  for i in range(length):
                      i_0 = input_0[i]
                      i_1 = input_1[length - 1 - i]
                      out = i_0 + (i_1 * 2) + s_0 + (s_1 * 2) + f_0
                      out = out * i * i_0 * i_1
                      s_0 = (s_0 + out) * 1.05
                      s_1 = (s_1 - out * 0.5) * 0.95
                      output_0.append(out)
                      output_1.append(f_0)
                      output_2.append(out * 1.5)
                  return output_0, output_1, output_2, s_0, s_1, s_2
              """
  
              def step(loop, free):
                  (i, s_0, s_1, s_2), (sc_0, sc_1, f_0, _) = loop, free
                  F = mx.sym if isinstance(i, mx.sym.Symbol) else mx.nd
                  i_0 = sc_0.take(i).squeeze(axis=0)
                  i_1 = sc_1.take(length - 1 - i).squeeze(axis=0)
                  out_0 = step_func(i_0, i_1, s_0, s_1, f_0)
                  out_0 = out_0 * i.reshape(
                      [1] * len(single_shape)
                  ).broadcast_to(single_shape)
                  out_1 = step_func(i_1, s_0, f_0, s_1, i_0)
                  out_1 = out_1 * i.reshape(
                      [1] * len(single_shape)
                  ).broadcast_to(single_shape)
                  return (
                      [F.dot(out_0, s_2), f_0, F.dot(s_2, out_1) * 1.5],
                      [
                          i + 1,
                          (s_0 + out_1) * 1.05,
                          (s_1 - out_0 * 0.5) * 0.95,
                          s_2,
                      ],
                  )
  
              return step
  
          case_id = 0
          for is_train in [True, False]:
              for step_func in step_funcs:
                  case_id += 1
                  _verify_while_loop(
                      func=make_func(step_func),
                      max_iterations=1000,
                      is_train=is_train,
                      is_for=True,
                      **params,
                  )
  
      # Case 0: the simpest case
      case_0()
      # Case 1.1.*
      case_1(
          cond=make_true_cond(),
          loop_var_shapes=[
              (1,),  # s
          ],
          free_var_shapes=[
              (1,),  # a
              (1,),  # b
          ],
          max_iterations=5,
          n_steps=5,
      )
      # Case 1.2.*
      case_1(
          cond=make_true_cond(),
          loop_var_shapes=[
              (2, 3, 4),  # s
          ],
          free_var_shapes=[
              (2, 3, 4),  # a
              (2, 3, 4),  # b
          ],
          max_iterations=3,
          n_steps=3,
      )
      # Case 1.3.*
      case_1(
          cond=make_false_cond(),
          loop_var_shapes=[
              (2, 3, 4),  # s
          ],
          free_var_shapes=[
              (2, 3, 4),  # a
              (2, 3, 4),  # b
          ],
          max_iterations=20,
          n_steps=0,
      )
      # Case 2.1.*
      case_2(
          cond=make_for_cond(length=5),
          loop_var_shapes=[
              (1,),  # i
              (2,),  # s
          ],
          free_var_shapes=[
              (100, 2),  # scanned
              (2,),  # f_1
              (3, 4, 5, 6),  # f_2, unused
          ],
          n_steps=5,
      )
      # Case 2.2.*
      case_2(
          cond=make_for_cond(length=3),
          loop_var_shapes=[
              (1,),  # i
              (2,),  # s
          ],
          free_var_shapes=[
              (30, 2),  # scanned
              (2,),  # f_1
              (3, 4, 5, 6),  # f_2, unused
          ],
          n_steps=3,
      )
      # Case 3.*
      case_3(
          length=5,
          cond=make_for_cond(length=5),
          loop_var_shapes=[
              (1,),  # i
              (2,),  # s_0
              (2,),  # s_1
          ],
          free_var_shapes=[
              (30, 2),  # sc_0
              (30, 2),  # sc_1
              (2,),  # f_0
              (3, 4, 5, 6),  # f_1, unused
          ],
          n_steps=5,
      )
      # Case 4.1.*
      case_4(
          length=4,
          cond=make_for_cond(length=4),
          single_shape=[5],
          loop_var_shapes=[
              (1,),  # i
              (5,),  # s_0
              (5,),  # s_1
              (23, 6, 8),  # s_2
          ],
          free_var_shapes=[
              (30, 5),  # sc_0
              (30, 5),  # sc_1
              (5,),  # f_0
              (3, 4, 5, 6),  # f_1, unused
          ],
          n_steps=4,
      )
      # Case 4.2.*
      case_4(
          length=5,
          cond=make_for_cond(length=5),
          single_shape=[5, 12],
          loop_var_shapes=[
              (1,),  # i
              (5, 12),  # s_0
              (5, 12),  # s_1
              (23, 6, 8),  # s_2
          ],
          free_var_shapes=[
              (30, 5, 12),  # sc_0
              (30, 5, 12),  # sc_1
              (5, 12),  # f_0
              (3, 4, 5, 6),  # f_1, unused
          ],
          n_steps=5,
      )
      # Case 5.1.*
      case_5(
          length=4,
          cond=make_for_cond(length=4),
          single_shape=[5],
          loop_var_shapes=[
              (1,),  # i
              (5,),  # s_0
              (5,),  # s_1
              (23, 6, 8),  # s_2
          ],
          free_var_shapes=[
              (30, 5),  # sc_0
              (30, 5),  # sc_1
              (5,),  # f_0
              (3, 4, 5, 6),  # f_1, unused
          ],
          n_steps=4,
      )
      # Case 5.2.*
      case_5(
          length=5,
          cond=make_for_cond(length=5),
          single_shape=[3, 4, 2],
          loop_var_shapes=[
              (1,),  # i
              (3, 4, 2),  # s_0
              (3, 4, 2),  # s_1
              (23, 6, 8),  # s_2
          ],
          free_var_shapes=[
              (30, 3, 4, 2),  # sc_0
              (30, 3, 4, 2),  # sc_1
              (3, 4, 2),  # f_0
              (3, 4, 5, 6),  # f_1, unused
          ],
          n_steps=5,
      )
      # Case 6.*
      case_6(
          length=5,
          cond=make_for_cond(length=5),
          single_shape=[5, 3],
          loop_var_shapes=[
              (1,),  # i
              (5, 3),  # s_0
              (5, 3),  # s_1
              (3, 5),  # s_2
          ],
          free_var_shapes=[
              (30, 5, 3),  # sc_0
              (30, 5, 3),  # sc_1
              (5, 3),  # f_0
              (3, 4, 5, 6),  # f_1, unused
          ],
          n_steps=5,
      )
  
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 2 for calling a global function with example usage: test_while_loop_for_foreach() and returns: function, function, function, in_,  unknown, step, step, step, step, step, step, out,  out, out,  unknown,  out, list,  list, list,  list, list,  list, list,  list