TestBatched.test_while.expect

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