TestBatched.test_while.expect

graph(%a.1_data : Dynamic
      %a.1_mask : Dynamic
      %a.1_dims : Dynamic
      %b_data : Dynamic
      %b_mask : Dynamic
      %b_dims : Dynamic) {
  %6 : int = prim::Constant[value=1]()
  %7 : int = prim::Constant[value=2147483647]()
  %8 : Dynamic = aten::gt(%a.1_data, %b_data)
  %9 : Dynamic = aten::mul(%a.1_mask, %b_mask)
  %10 : Dynamic = aten::__or__(%a.1_dims, %b_dims)
  %11 : bool = prim::TensorToBool(%8)
  %12 : int = prim::Constant[value=0]()
  %13 : Dynamic = aten::mul(%8, %9)
  %14 : Dynamic = aten::sum(%13)
  %15 : Dynamic = aten::gt(%14, %12)
  %16 : bool = prim::TensorToBool(%15)
  %17 : Dynamic, %18 : Dynamic, %19 : Dynamic, %a : Dynamic, %21 : Dynamic, %22 : Dynamic = prim::Loop(%7, %16, %8, %9, %10, %a.1_data, %a.1_mask, %a.1_dims)
    block0(%loop_num : int, %cond_data.2 : Dynamic, %cond_mask.3 : Dynamic, %cond_dims : Dynamic, %6_data : Dynamic, %6_mask : Dynamic, %6_dims : Dynamic) {
      %30 : Long() = prim::NumToTensor(%6)
      %alpha : float = prim::TensorToNum(%30)
      %data.1 : Dynamic = aten::sub(%6_data, %b_data, %alpha)
      %mask : Dynamic = aten::mul(%6_mask, %b_mask)
      %dims : Dynamic = aten::__or__(%6_dims, %b_dims)
      %35 : Dynamic = aten::gt(%data.1, %b_data)
      %36 : Dynamic = aten::mul(%mask, %b_mask)
      %37 : Dynamic = aten::__or__(%dims, %b_dims)
      %38 : bool = prim::TensorToBool(%35)
      %39 : bool = prim::Constant[value=1]()
      %40 : int = prim::Constant[value=1]()
      %41 : Dynamic = aten::type_as(%cond_mask.3, %cond_data.2)
      %cond_mask.1 : Dynamic = aten::mul(%cond_data.2, %41)
      %43 : int = aten::dim(%cond_mask.1)
      %44 : bool = aten::eq(%43, %40)
      %cond_data : Dynamic, %cond_mask : Dynamic, %data : Dynamic = prim::If(%44)
        block0() {
          %48 : int = aten::dim(%data.1)
          %49 : int = aten::sub(%48, %40)
          %data.3 : Dynamic = prim::Loop(%49, %39, %cond_mask.1)
            block0(%_ : int, %52 : Dynamic) {
              %53 : int = aten::dim(%52)
              %data.2 : Dynamic = aten::unsqueeze(%52, %53)
              -> (%39, %data.2)
            }
          %cond_data.1 : Dynamic = aten::expand_as(%data.3, %data.1)
          %cond_mask.2 : Dynamic = aten::expand_as(%data.3, %mask)
          -> (%cond_data.1, %cond_mask.2, %data.3)
        }
        block1() {
          -> (%cond_mask.1, %cond_mask.1, %cond_mask.1)
        }
      %res_data : Dynamic = aten::where(%cond_data, %data.1, %6_data)
      %res_mask : Dynamic = aten::where(%cond_mask, %mask, %6_mask)
      %res_dims : Dynamic = aten::__or__(%dims, %6_dims)
      %60 : int = prim::Constant[value=0]()
      %61 : Dynamic = aten::mul(%35, %36)
      %62 : Dynamic = aten::sum(%61)
      %63 : Dynamic = aten::gt(%62, %60)
      %64 : bool = prim::TensorToBool(%63)
      -> (%64, %35, %36, %37, %res_data, %res_mask, %res_dims)
    }
  return (%a, %21, %22);
}