Skip to content

Instantly share code, notes, and snippets.

@shunting314

shunting314/pyhpc_isoneutral_mixing

Last active March 10, 2022 01:34
Show Gist options
  • Save shunting314/d3f09b92196654303f0c9454db83b322 to your computer and use it in GitHub Desktop.
Save shunting314/d3f09b92196654303f0c9454db83b322 to your computer and use it in GitHub Desktop.
Download ZIP
An fx graph captured for pyhpc_isoneutral_mixing by torchdynamo
Raw
pyhpc_isoneutral_mixing
def forward(self, maskT : torch.Tensor, maskU : torch.Tensor, maskV : torch.Tensor, maskW : torch.Tensor, dxt : torch.Tensor, dxu : torch.Tensor, dyt : torch.Tensor, dyu : torch.Tensor, dzt : torch.Tensor, dzw : torch.Tensor, cost : torch.Tensor, cosu : torch.Tensor, salt : torch.Tensor, temp : torch.Tensor, zt : torch.Tensor, K_iso : torch.Tensor, K_11 : torch.Tensor, K_22 : torch.Tensor, K_33 : torch.Tensor, Ai_ez : torch.Tensor, Ai_nz : torch.Tensor, Ai_bx : torch.Tensor, Ai_by : torch.Tensor):
mask_t = maskT
mask_u = maskU
mask_v = maskV
mask_w = maskW
k_iso = K_iso
k_11 = K_11
k_22 = K_22
k_33 = K_33
ai_ez = Ai_ez
ai_nz = Ai_nz
ai_bx = Ai_bx
ai_by = Ai_by
zeros_like = torch.zeros_like(k_11)
zeros_like_1 = torch.zeros_like(k_11)
zeros_like_2 = torch.zeros_like(k_11)
zeros_like_3 = torch.zeros_like(k_11)
zeros_like_4 = torch.zeros_like(k_11)
zeros_like_5 = torch.zeros_like(k_11)
getitem_1 = temp[(slice(None, None, None), slice(None, None, None), slice(None, None, None), 0)]
abs_1 = torch.abs(zt)
neg = -abs_1; abs_1 = None
sub = neg - 0.0; neg = None
sub_1 = getitem_1 - 9.850000000000023; getitem_1 = None
mul = 1e-05 * sub_1; sub_1 = None
mul_1 = 1.0790999999999999e-07 * sub; sub = None
mul_2 = mul_1 * 1024.0; mul_1 = None
sub_2 = 1 - mul_2; mul_2 = None
mul_3 = 0.000167 * sub_2; sub_2 = None
add = mul + mul_3; mul = mul_3 = None
neg_1 = -add; add = None
mul_4 = neg_1 * 1024.0; neg_1 = None
mul_5 = mask_t * mul_4; mul_4 = None
getitem_3 = temp[(slice(None, None, None), slice(None, None, None), slice(None, None, None), 0)]
abs_2 = torch.abs(zt); zt = None
ones_like = torch.ones_like(getitem_3); getitem_3 = None
mul_6 = 0.79872 * ones_like; ones_like = None
mul_7 = mask_t * mul_6; mask_t = mul_6 = None
getitem_4 = mask_w[(slice(None, None, None), slice(None, None, None), slice(None, -1, None))]
getitem_5 = temp[(slice(None, None, None), slice(None, None, None), slice(1, None, None), 0)]
getitem_6 = temp[(slice(None, None, None), slice(None, None, None), slice(None, -1, None), 0)]
sub_3 = getitem_5 - getitem_6; getitem_5 = getitem_6 = None
mul_8 = getitem_4 * sub_3; getitem_4 = sub_3 = None
getitem_7 = dzw[(None, None, slice(None, -1, None))]
truediv = mul_8 / getitem_7; mul_8 = getitem_7 = None
zeros_like_4[(slice(None, None, None), slice(None, None, None), slice(None, -1, None))] = truediv; setitem = zeros_like_4; truediv = None
getitem_8 = mask_w[(slice(None, None, None), slice(None, None, None), slice(None, -1, None))]
getitem_9 = salt[(slice(None, None, None), slice(None, None, None), slice(1, None, None), 0)]
getitem_10 = salt[(slice(None, None, None), slice(None, None, None), slice(None, -1, None), 0)]
sub_4 = getitem_9 - getitem_10; getitem_9 = getitem_10 = None
mul_9 = getitem_8 * sub_4; getitem_8 = sub_4 = None
getitem_11 = dzw[(None, None, slice(None, -1, None))]
truediv_1 = mul_9 / getitem_11; mul_9 = getitem_11 = None
zeros_like_5[(slice(None, None, None), slice(None, None, None), slice(None, -1, None))] = truediv_1; setitem_1 = zeros_like_5; truediv_1 = None
getitem_12 = mask_u[(slice(None, -1, None), slice(None, None, None), slice(None, None, None))]
getitem_13 = temp[(slice(1, None, None), slice(None, None, None), slice(None, None, None), 0)]
getitem_14 = temp[(slice(None, -1, None), slice(None, None, None), slice(None, None, None), 0)]
sub_5 = getitem_13 - getitem_14; getitem_13 = getitem_14 = None
mul_10 = getitem_12 * sub_5; getitem_12 = sub_5 = None
getitem_15 = dxu[(slice(None, -1, None), None, None)]
getitem_16 = cost[(None, slice(None, None, None), None)]
mul_11 = getitem_15 * getitem_16; getitem_15 = getitem_16 = None
truediv_2 = mul_10 / mul_11; mul_10 = mul_11 = None
zeros_like[(slice(None, -1, None), slice(None, None, None), slice(None, None, None))] = truediv_2; setitem_2 = zeros_like; truediv_2 = None
getitem_17 = mask_u[(slice(None, -1, None), slice(None, None, None), slice(None, None, None))]
getitem_18 = salt[(slice(1, None, None), slice(None, None, None), slice(None, None, None), 0)]
getitem_19 = salt[(slice(None, -1, None), slice(None, None, None), slice(None, None, None), 0)]
sub_6 = getitem_18 - getitem_19; getitem_18 = getitem_19 = None
mul_12 = getitem_17 * sub_6; getitem_17 = sub_6 = None
getitem_20 = dxu[(slice(None, -1, None), None, None)]
getitem_21 = cost[(None, slice(None, None, None), None)]
mul_13 = getitem_20 * getitem_21; getitem_20 = getitem_21 = None
truediv_3 = mul_12 / mul_13; mul_12 = mul_13 = None
zeros_like_1[(slice(None, -1, None), slice(None, None, None), slice(None, None, None))] = truediv_3; setitem_3 = zeros_like_1; truediv_3 = None
getitem_22 = mask_v[(slice(None, None, None), slice(None, -1, None), slice(None, None, None))]
getitem_23 = temp[(slice(None, None, None), slice(1, None, None), slice(None, None, None), 0)]
getitem_24 = temp[(slice(None, None, None), slice(None, -1, None), slice(None, None, None), 0)]; temp = None
sub_7 = getitem_23 - getitem_24; getitem_23 = getitem_24 = None
mul_14 = getitem_22 * sub_7; getitem_22 = sub_7 = None
getitem_25 = dyu[(None, slice(None, -1, None), None)]
truediv_4 = mul_14 / getitem_25; mul_14 = getitem_25 = None
zeros_like_2[(slice(None, None, None), slice(None, -1, None), slice(None, None, None))] = truediv_4; setitem_4 = zeros_like_2; truediv_4 = None
getitem_26 = mask_v[(slice(None, None, None), slice(None, -1, None), slice(None, None, None))]
getitem_27 = salt[(slice(None, None, None), slice(1, None, None), slice(None, None, None), 0)]
getitem_28 = salt[(slice(None, None, None), slice(None, -1, None), slice(None, None, None), 0)]; salt = None
sub_8 = getitem_27 - getitem_28; getitem_27 = getitem_28 = None
mul_15 = getitem_26 * sub_8; getitem_26 = sub_8 = None
getitem_29 = dyu[(None, slice(None, -1, None), None)]
truediv_5 = mul_15 / getitem_29; mul_15 = getitem_29 = None
zeros_like_3[(slice(None, None, None), slice(None, -1, None), slice(None, None, None))] = truediv_5; setitem_5 = zeros_like_3; truediv_5 = None
zeros_like_6 = torch.zeros_like(k_11)
getitem_30 = k_iso[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
getitem_31 = k_iso[(slice(1, -2, None), slice(2, -2, None), slice(None, -1, None))]
add_1 = getitem_30 + getitem_31; getitem_30 = getitem_31 = None
getitem_32 = k_iso[(slice(2, -1, None), slice(2, -2, None), slice(1, None, None))]
add_2 = add_1 + getitem_32; add_1 = getitem_32 = None
getitem_33 = k_iso[(slice(2, -1, None), slice(2, -2, None), slice(None, -1, None))]
add_3 = add_2 + getitem_33; add_2 = getitem_33 = None
mul_16 = 0.25 * add_3; add_3 = None
zeros_like_6[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))] = mul_16; setitem_6 = zeros_like_6; mul_16 = None
getitem_34 = k_iso[(slice(1, -2, None), slice(2, -2, None), 0)]
getitem_35 = k_iso[(slice(2, -1, None), slice(2, -2, None), 0)]
add_4 = getitem_34 + getitem_35; getitem_34 = getitem_35 = None
mul_17 = 0.5 * add_4; add_4 = None
zeros_like_6[(slice(1, -2, None), slice(2, -2, None), 0)] = mul_17; setitem_7 = zeros_like_6; mul_17 = None
zeros_like_7 = torch.zeros_like(k_11)
getitem_36 = zeros_like_7[(slice(1, -2, None), slice(2, -2, None))]; zeros_like_7 = None
getitem_37 = mul_5[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
getitem_38 = zeros_like[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_18 = getitem_37 * getitem_38; getitem_37 = getitem_38 = None
getitem_39 = mul_7[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
getitem_40 = zeros_like_1[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_19 = getitem_39 * getitem_40; getitem_39 = getitem_40 = None
add_5 = mul_18 + mul_19; mul_18 = mul_19 = None
getitem_41 = mul_5[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
getitem_42 = zeros_like_4[(slice(1, -2, None), slice(2, -2, None), slice(None, 25, None))]
mul_20 = getitem_41 * getitem_42; getitem_41 = getitem_42 = None
getitem_43 = mul_7[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
getitem_44 = zeros_like_5[(slice(1, -2, None), slice(2, -2, None), slice(None, 25, None))]
mul_21 = getitem_43 * getitem_44; getitem_43 = getitem_44 = None
add_6 = mul_20 + mul_21; mul_20 = mul_21 = None
neg_2 = -add_5; add_5 = None
tensor = torch.tensor([0.0], device = device(type='cuda', index=0))
min_1 = torch.min(add_6, tensor); add_6 = tensor = None
sub_9 = min_1 - 1e-20; min_1 = None
truediv_6 = neg_2 / sub_9; neg_2 = sub_9 = None
abs_3 = torch.abs(truediv_6)
neg_3 = -abs_3; abs_3 = None
add_7 = neg_3 + 0.001; neg_3 = None
truediv_7 = add_7 / 0.001; add_7 = None
tanh = torch.tanh(truediv_7); truediv_7 = None
add_8 = 1.0 + tanh; tanh = None
mul_22 = 0.5 * add_8; add_8 = None
getitem_45 = getitem_36[(slice(None, None, None), slice(None, None, None), slice(1, None, None))]
getitem_46 = dzw[(None, None, slice(None, 25, None))]
getitem_47 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_23 = getitem_46 * getitem_47; getitem_46 = getitem_47 = None
tensor_1 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_48 = zeros_like_6[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_24 = getitem_48 * mul_22; getitem_48 = None
max_1 = torch.max(tensor_1, mul_24); tensor_1 = mul_24 = None
mul_25 = mul_23 * max_1; mul_23 = max_1 = None
getitem_45 += mul_25; iadd = getitem_45; getitem_45 = mul_25 = None
getitem_36[(slice(None, None, None), slice(None, None, None), slice(1, None, None))] = iadd; setitem_8 = getitem_36; iadd = None
mul_26 = mul_22 * truediv_6; mul_22 = truediv_6 = None
getitem_49 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_27 = mul_26 * getitem_49; mul_26 = getitem_49 = None
ai_ez[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None), 0, 0)] = mul_27; setitem_9 = ai_ez; mul_27 = None
getitem_50 = mul_5[(slice(2, -1, None), slice(2, -2, None), slice(1, None, None))]
getitem_51 = zeros_like[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_28 = getitem_50 * getitem_51; getitem_50 = getitem_51 = None
getitem_52 = mul_7[(slice(2, -1, None), slice(2, -2, None), slice(1, None, None))]
getitem_53 = zeros_like_1[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_29 = getitem_52 * getitem_53; getitem_52 = getitem_53 = None
add_9 = mul_28 + mul_29; mul_28 = mul_29 = None
getitem_54 = mul_5[(slice(2, -1, None), slice(2, -2, None), slice(1, None, None))]
getitem_55 = zeros_like_4[(slice(2, -1, None), slice(2, -2, None), slice(None, 25, None))]
mul_30 = getitem_54 * getitem_55; getitem_54 = getitem_55 = None
getitem_56 = mul_7[(slice(2, -1, None), slice(2, -2, None), slice(1, None, None))]
getitem_57 = zeros_like_5[(slice(2, -1, None), slice(2, -2, None), slice(None, 25, None))]
mul_31 = getitem_56 * getitem_57; getitem_56 = getitem_57 = None
add_10 = mul_30 + mul_31; mul_30 = mul_31 = None
neg_4 = -add_9; add_9 = None
tensor_2 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_2 = torch.min(add_10, tensor_2); add_10 = tensor_2 = None
sub_10 = min_2 - 1e-20; min_2 = None
truediv_8 = neg_4 / sub_10; neg_4 = sub_10 = None
abs_4 = torch.abs(truediv_8)
neg_5 = -abs_4; abs_4 = None
add_11 = neg_5 + 0.001; neg_5 = None
truediv_9 = add_11 / 0.001; add_11 = None
tanh_1 = torch.tanh(truediv_9); truediv_9 = None
add_12 = 1.0 + tanh_1; tanh_1 = None
mul_32 = 0.5 * add_12; add_12 = None
getitem_58 = getitem_36[(slice(None, None, None), slice(None, None, None), slice(1, None, None))]
getitem_59 = dzw[(None, None, slice(None, 25, None))]
getitem_60 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_33 = getitem_59 * getitem_60; getitem_59 = getitem_60 = None
tensor_3 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_61 = zeros_like_6[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_34 = getitem_61 * mul_32; getitem_61 = None
max_2 = torch.max(tensor_3, mul_34); tensor_3 = mul_34 = None
mul_35 = mul_33 * max_2; mul_33 = max_2 = None
getitem_58 += mul_35; iadd_1 = getitem_58; getitem_58 = mul_35 = None
getitem_36[(slice(None, None, None), slice(None, None, None), slice(1, None, None))] = iadd_1; setitem_10 = getitem_36; iadd_1 = None
mul_36 = mul_32 * truediv_8; mul_32 = truediv_8 = None
getitem_62 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None))]
mul_37 = mul_36 * getitem_62; mul_36 = getitem_62 = None
ai_ez[(slice(1, -2, None), slice(2, -2, None), slice(1, None, None), 1, 0)] = mul_37; setitem_11 = ai_ez; mul_37 = None
getitem_63 = mul_5[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
getitem_64 = zeros_like[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_38 = getitem_63 * getitem_64; getitem_63 = getitem_64 = None
getitem_65 = mul_7[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
getitem_66 = zeros_like_1[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_39 = getitem_65 * getitem_66; getitem_65 = getitem_66 = None
add_13 = mul_38 + mul_39; mul_38 = mul_39 = None
getitem_67 = mul_5[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
getitem_68 = zeros_like_4[(slice(1, -2, None), slice(2, -2, None), slice(None, 26, None))]
mul_40 = getitem_67 * getitem_68; getitem_67 = getitem_68 = None
getitem_69 = mul_7[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
getitem_70 = zeros_like_5[(slice(1, -2, None), slice(2, -2, None), slice(None, 26, None))]
mul_41 = getitem_69 * getitem_70; getitem_69 = getitem_70 = None
add_14 = mul_40 + mul_41; mul_40 = mul_41 = None
neg_6 = -add_13; add_13 = None
tensor_4 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_3 = torch.min(add_14, tensor_4); add_14 = tensor_4 = None
sub_11 = min_3 - 1e-20; min_3 = None
truediv_10 = neg_6 / sub_11; neg_6 = sub_11 = None
abs_5 = torch.abs(truediv_10)
neg_7 = -abs_5; abs_5 = None
add_15 = neg_7 + 0.001; neg_7 = None
truediv_11 = add_15 / 0.001; add_15 = None
tanh_2 = torch.tanh(truediv_11); truediv_11 = None
add_16 = 1.0 + tanh_2; tanh_2 = None
mul_42 = 0.5 * add_16; add_16 = None
getitem_71 = getitem_36[(slice(None, None, None), slice(None, None, None), slice(0, None, None))]
getitem_72 = dzw[(None, None, slice(None, 26, None))]
getitem_73 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_43 = getitem_72 * getitem_73; getitem_72 = getitem_73 = None
tensor_5 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_74 = zeros_like_6[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_44 = getitem_74 * mul_42; getitem_74 = None
max_3 = torch.max(tensor_5, mul_44); tensor_5 = mul_44 = None
mul_45 = mul_43 * max_3; mul_43 = max_3 = None
getitem_71 += mul_45; iadd_2 = getitem_71; getitem_71 = mul_45 = None
getitem_36[(slice(None, None, None), slice(None, None, None), slice(0, None, None))] = iadd_2; setitem_12 = getitem_36; iadd_2 = None
mul_46 = mul_42 * truediv_10; mul_42 = truediv_10 = None
getitem_75 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_47 = mul_46 * getitem_75; mul_46 = getitem_75 = None
ai_ez[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None), 0, 1)] = mul_47; setitem_13 = ai_ez; mul_47 = None
getitem_76 = mul_5[(slice(2, -1, None), slice(2, -2, None), slice(0, None, None))]
getitem_77 = zeros_like[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_48 = getitem_76 * getitem_77; getitem_76 = getitem_77 = None
getitem_78 = mul_7[(slice(2, -1, None), slice(2, -2, None), slice(0, None, None))]
getitem_79 = zeros_like_1[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_49 = getitem_78 * getitem_79; getitem_78 = getitem_79 = None
add_17 = mul_48 + mul_49; mul_48 = mul_49 = None
getitem_80 = mul_5[(slice(2, -1, None), slice(2, -2, None), slice(0, None, None))]
getitem_81 = zeros_like_4[(slice(2, -1, None), slice(2, -2, None), slice(None, 26, None))]
mul_50 = getitem_80 * getitem_81; getitem_80 = getitem_81 = None
getitem_82 = mul_7[(slice(2, -1, None), slice(2, -2, None), slice(0, None, None))]
getitem_83 = zeros_like_5[(slice(2, -1, None), slice(2, -2, None), slice(None, 26, None))]
mul_51 = getitem_82 * getitem_83; getitem_82 = getitem_83 = None
add_18 = mul_50 + mul_51; mul_50 = mul_51 = None
neg_8 = -add_17; add_17 = None
tensor_6 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_4 = torch.min(add_18, tensor_6); add_18 = tensor_6 = None
sub_12 = min_4 - 1e-20; min_4 = None
truediv_12 = neg_8 / sub_12; neg_8 = sub_12 = None
abs_6 = torch.abs(truediv_12)
neg_9 = -abs_6; abs_6 = None
add_19 = neg_9 + 0.001; neg_9 = None
truediv_13 = add_19 / 0.001; add_19 = None
tanh_3 = torch.tanh(truediv_13); truediv_13 = None
add_20 = 1.0 + tanh_3; tanh_3 = None
mul_52 = 0.5 * add_20; add_20 = None
getitem_84 = getitem_36[(slice(None, None, None), slice(None, None, None), slice(0, None, None))]
getitem_85 = dzw[(None, None, slice(None, 26, None))]
getitem_86 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_53 = getitem_85 * getitem_86; getitem_85 = getitem_86 = None
tensor_7 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_87 = zeros_like_6[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]
mul_54 = getitem_87 * mul_52; getitem_87 = None
max_4 = torch.max(tensor_7, mul_54); tensor_7 = mul_54 = None
mul_55 = mul_53 * max_4; mul_53 = max_4 = None
getitem_84 += mul_55; iadd_3 = getitem_84; getitem_84 = mul_55 = None
getitem_36[(slice(None, None, None), slice(None, None, None), slice(0, None, None))] = iadd_3; setitem_14 = getitem_36; iadd_3 = None
mul_56 = mul_52 * truediv_12; mul_52 = truediv_12 = None
getitem_88 = mask_u[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None))]; mask_u = None
mul_57 = mul_56 * getitem_88; mul_56 = getitem_88 = None
ai_ez[(slice(1, -2, None), slice(2, -2, None), slice(0, None, None), 1, 1)] = mul_57; setitem_15 = ai_ez; ai_ez = mul_57 = None
getitem_89 = dzt[(None, None, slice(None, None, None))]
mul_58 = 4.0 * getitem_89; getitem_89 = None
truediv_14 = getitem_36 / mul_58; getitem_36 = mul_58 = None
k_11[(slice(1, -2, None), slice(2, -2, None), slice(None, None, None))] = truediv_14; setitem_16 = k_11; truediv_14 = None
zeros_like_6[Ellipsis] = 0; setitem_17 = zeros_like_6
getitem_90 = k_iso[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
getitem_91 = k_iso[(slice(2, -2, None), slice(1, -2, None), slice(None, -1, None))]
add_21 = getitem_90 + getitem_91; getitem_90 = getitem_91 = None
getitem_92 = k_iso[(slice(2, -2, None), slice(2, -1, None), slice(1, None, None))]
add_22 = add_21 + getitem_92; add_21 = getitem_92 = None
getitem_93 = k_iso[(slice(2, -2, None), slice(2, -1, None), slice(None, -1, None))]
add_23 = add_22 + getitem_93; add_22 = getitem_93 = None
mul_59 = 0.25 * add_23; add_23 = None
zeros_like_6[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))] = mul_59; setitem_18 = zeros_like_6; mul_59 = None
getitem_94 = k_iso[(slice(2, -2, None), slice(1, -2, None), 0)]
getitem_95 = k_iso[(slice(2, -2, None), slice(2, -1, None), 0)]
add_24 = getitem_94 + getitem_95; getitem_94 = getitem_95 = None
mul_60 = 0.5 * add_24; add_24 = None
zeros_like_6[(slice(2, -2, None), slice(1, -2, None), 0)] = mul_60; setitem_19 = zeros_like_6; mul_60 = None
zeros_like_8 = torch.zeros_like(k_11)
getitem_96 = zeros_like_8[(slice(2, -2, None), slice(1, -2, None))]; zeros_like_8 = None
getitem_97 = mul_5[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
getitem_98 = zeros_like_2[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_61 = getitem_97 * getitem_98; getitem_97 = getitem_98 = None
getitem_99 = mul_7[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
getitem_100 = zeros_like_3[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_62 = getitem_99 * getitem_100; getitem_99 = getitem_100 = None
add_25 = mul_61 + mul_62; mul_61 = mul_62 = None
getitem_101 = mul_5[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
getitem_102 = zeros_like_4[(slice(2, -2, None), slice(1, -2, None), slice(None, 25, None))]
mul_63 = getitem_101 * getitem_102; getitem_101 = getitem_102 = None
getitem_103 = mul_7[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
getitem_104 = zeros_like_5[(slice(2, -2, None), slice(1, -2, None), slice(None, 25, None))]
mul_64 = getitem_103 * getitem_104; getitem_103 = getitem_104 = None
add_26 = mul_63 + mul_64; mul_63 = mul_64 = None
neg_10 = -add_25; add_25 = None
tensor_8 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_5 = torch.min(tensor_8, add_26); tensor_8 = add_26 = None
sub_13 = min_5 - 1e-20; min_5 = None
truediv_15 = neg_10 / sub_13; neg_10 = sub_13 = None
abs_7 = torch.abs(truediv_15)
neg_11 = -abs_7; abs_7 = None
add_27 = neg_11 + 0.001; neg_11 = None
truediv_16 = add_27 / 0.001; add_27 = None
tanh_4 = torch.tanh(truediv_16); truediv_16 = None
add_28 = 1.0 + tanh_4; tanh_4 = None
mul_65 = 0.5 * add_28; add_28 = None
getitem_105 = getitem_96[(slice(None, None, None), slice(None, None, None), slice(1, None, None))]
getitem_106 = dzw[(None, None, slice(None, 25, None))]
getitem_107 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_66 = getitem_106 * getitem_107; getitem_106 = getitem_107 = None
tensor_9 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_108 = zeros_like_6[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_67 = getitem_108 * mul_65; getitem_108 = None
max_5 = torch.max(tensor_9, mul_67); tensor_9 = mul_67 = None
mul_68 = mul_66 * max_5; mul_66 = max_5 = None
getitem_105 += mul_68; iadd_4 = getitem_105; getitem_105 = mul_68 = None
getitem_96[(slice(None, None, None), slice(None, None, None), slice(1, None, None))] = iadd_4; setitem_20 = getitem_96; iadd_4 = None
mul_69 = mul_65 * truediv_15; mul_65 = truediv_15 = None
getitem_109 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_70 = mul_69 * getitem_109; mul_69 = getitem_109 = None
ai_nz[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None), 0, 0)] = mul_70; setitem_21 = ai_nz; mul_70 = None
getitem_110 = mul_5[(slice(2, -2, None), slice(2, -1, None), slice(1, None, None))]
getitem_111 = zeros_like_2[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_71 = getitem_110 * getitem_111; getitem_110 = getitem_111 = None
getitem_112 = mul_7[(slice(2, -2, None), slice(2, -1, None), slice(1, None, None))]
getitem_113 = zeros_like_3[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_72 = getitem_112 * getitem_113; getitem_112 = getitem_113 = None
add_29 = mul_71 + mul_72; mul_71 = mul_72 = None
getitem_114 = mul_5[(slice(2, -2, None), slice(2, -1, None), slice(1, None, None))]
getitem_115 = zeros_like_4[(slice(2, -2, None), slice(2, -1, None), slice(None, 25, None))]
mul_73 = getitem_114 * getitem_115; getitem_114 = getitem_115 = None
getitem_116 = mul_7[(slice(2, -2, None), slice(2, -1, None), slice(1, None, None))]
getitem_117 = zeros_like_5[(slice(2, -2, None), slice(2, -1, None), slice(None, 25, None))]
mul_74 = getitem_116 * getitem_117; getitem_116 = getitem_117 = None
add_30 = mul_73 + mul_74; mul_73 = mul_74 = None
neg_12 = -add_29; add_29 = None
tensor_10 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_6 = torch.min(tensor_10, add_30); tensor_10 = add_30 = None
sub_14 = min_6 - 1e-20; min_6 = None
truediv_17 = neg_12 / sub_14; neg_12 = sub_14 = None
abs_8 = torch.abs(truediv_17)
neg_13 = -abs_8; abs_8 = None
add_31 = neg_13 + 0.001; neg_13 = None
truediv_18 = add_31 / 0.001; add_31 = None
tanh_5 = torch.tanh(truediv_18); truediv_18 = None
add_32 = 1.0 + tanh_5; tanh_5 = None
mul_75 = 0.5 * add_32; add_32 = None
getitem_118 = getitem_96[(slice(None, None, None), slice(None, None, None), slice(1, None, None))]
getitem_119 = dzw[(None, None, slice(None, 25, None))]
getitem_120 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_76 = getitem_119 * getitem_120; getitem_119 = getitem_120 = None
tensor_11 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_121 = zeros_like_6[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_77 = getitem_121 * mul_75; getitem_121 = None
max_6 = torch.max(tensor_11, mul_77); tensor_11 = mul_77 = None
mul_78 = mul_76 * max_6; mul_76 = max_6 = None
getitem_118 += mul_78; iadd_5 = getitem_118; getitem_118 = mul_78 = None
getitem_96[(slice(None, None, None), slice(None, None, None), slice(1, None, None))] = iadd_5; setitem_22 = getitem_96; iadd_5 = None
mul_79 = mul_75 * truediv_17; mul_75 = truediv_17 = None
getitem_122 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None))]
mul_80 = mul_79 * getitem_122; mul_79 = getitem_122 = None
ai_nz[(slice(2, -2, None), slice(1, -2, None), slice(1, None, None), 1, 0)] = mul_80; setitem_23 = ai_nz; mul_80 = None
getitem_123 = mul_5[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
getitem_124 = zeros_like_2[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_81 = getitem_123 * getitem_124; getitem_123 = getitem_124 = None
getitem_125 = mul_7[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
getitem_126 = zeros_like_3[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_82 = getitem_125 * getitem_126; getitem_125 = getitem_126 = None
add_33 = mul_81 + mul_82; mul_81 = mul_82 = None
getitem_127 = mul_5[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
getitem_128 = zeros_like_4[(slice(2, -2, None), slice(1, -2, None), slice(None, 26, None))]
mul_83 = getitem_127 * getitem_128; getitem_127 = getitem_128 = None
getitem_129 = mul_7[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
getitem_130 = zeros_like_5[(slice(2, -2, None), slice(1, -2, None), slice(None, 26, None))]
mul_84 = getitem_129 * getitem_130; getitem_129 = getitem_130 = None
add_34 = mul_83 + mul_84; mul_83 = mul_84 = None
neg_14 = -add_33; add_33 = None
tensor_12 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_7 = torch.min(tensor_12, add_34); tensor_12 = add_34 = None
sub_15 = min_7 - 1e-20; min_7 = None
truediv_19 = neg_14 / sub_15; neg_14 = sub_15 = None
abs_9 = torch.abs(truediv_19)
neg_15 = -abs_9; abs_9 = None
add_35 = neg_15 + 0.001; neg_15 = None
truediv_20 = add_35 / 0.001; add_35 = None
tanh_6 = torch.tanh(truediv_20); truediv_20 = None
add_36 = 1.0 + tanh_6; tanh_6 = None
mul_85 = 0.5 * add_36; add_36 = None
getitem_131 = getitem_96[(slice(None, None, None), slice(None, None, None), slice(0, None, None))]
getitem_132 = dzw[(None, None, slice(None, 26, None))]
getitem_133 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_86 = getitem_132 * getitem_133; getitem_132 = getitem_133 = None
tensor_13 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_134 = zeros_like_6[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_87 = getitem_134 * mul_85; getitem_134 = None
max_7 = torch.max(tensor_13, mul_87); tensor_13 = mul_87 = None
mul_88 = mul_86 * max_7; mul_86 = max_7 = None
getitem_131 += mul_88; iadd_6 = getitem_131; getitem_131 = mul_88 = None
getitem_96[(slice(None, None, None), slice(None, None, None), slice(0, None, None))] = iadd_6; setitem_24 = getitem_96; iadd_6 = None
mul_89 = mul_85 * truediv_19; mul_85 = truediv_19 = None
getitem_135 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_90 = mul_89 * getitem_135; mul_89 = getitem_135 = None
ai_nz[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None), 0, 1)] = mul_90; setitem_25 = ai_nz; mul_90 = None
getitem_136 = mul_5[(slice(2, -2, None), slice(2, -1, None), slice(0, None, None))]
getitem_137 = zeros_like_2[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_91 = getitem_136 * getitem_137; getitem_136 = getitem_137 = None
getitem_138 = mul_7[(slice(2, -2, None), slice(2, -1, None), slice(0, None, None))]
getitem_139 = zeros_like_3[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_92 = getitem_138 * getitem_139; getitem_138 = getitem_139 = None
add_37 = mul_91 + mul_92; mul_91 = mul_92 = None
getitem_140 = mul_5[(slice(2, -2, None), slice(2, -1, None), slice(0, None, None))]
getitem_141 = zeros_like_4[(slice(2, -2, None), slice(2, -1, None), slice(None, 26, None))]
mul_93 = getitem_140 * getitem_141; getitem_140 = getitem_141 = None
getitem_142 = mul_7[(slice(2, -2, None), slice(2, -1, None), slice(0, None, None))]
getitem_143 = zeros_like_5[(slice(2, -2, None), slice(2, -1, None), slice(None, 26, None))]
mul_94 = getitem_142 * getitem_143; getitem_142 = getitem_143 = None
add_38 = mul_93 + mul_94; mul_93 = mul_94 = None
neg_16 = -add_37; add_37 = None
tensor_14 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_8 = torch.min(tensor_14, add_38); tensor_14 = add_38 = None
sub_16 = min_8 - 1e-20; min_8 = None
truediv_21 = neg_16 / sub_16; neg_16 = sub_16 = None
abs_10 = torch.abs(truediv_21)
neg_17 = -abs_10; abs_10 = None
add_39 = neg_17 + 0.001; neg_17 = None
truediv_22 = add_39 / 0.001; add_39 = None
tanh_7 = torch.tanh(truediv_22); truediv_22 = None
add_40 = 1.0 + tanh_7; tanh_7 = None
mul_95 = 0.5 * add_40; add_40 = None
getitem_144 = getitem_96[(slice(None, None, None), slice(None, None, None), slice(0, None, None))]
getitem_145 = dzw[(None, None, slice(None, 26, None))]; dzw = None
getitem_146 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]
mul_96 = getitem_145 * getitem_146; getitem_145 = getitem_146 = None
tensor_15 = torch.tensor([50.0], device = device(type='cuda', index=0))
getitem_147 = zeros_like_6[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]; zeros_like_6 = None
mul_97 = getitem_147 * mul_95; getitem_147 = None
max_8 = torch.max(tensor_15, mul_97); tensor_15 = mul_97 = None
mul_98 = mul_96 * max_8; mul_96 = max_8 = None
getitem_144 += mul_98; iadd_7 = getitem_144; getitem_144 = mul_98 = None
getitem_96[(slice(None, None, None), slice(None, None, None), slice(0, None, None))] = iadd_7; setitem_26 = getitem_96; iadd_7 = None
mul_99 = mul_95 * truediv_21; mul_95 = truediv_21 = None
getitem_148 = mask_v[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None))]; mask_v = None
mul_100 = mul_99 * getitem_148; mul_99 = getitem_148 = None
ai_nz[(slice(2, -2, None), slice(1, -2, None), slice(0, None, None), 1, 1)] = mul_100; setitem_27 = ai_nz; ai_nz = mul_100 = None
getitem_149 = dzt[(None, None, slice(None, None, None))]; dzt = None
mul_101 = 4.0 * getitem_149; getitem_149 = None
truediv_23 = getitem_96 / mul_101; getitem_96 = mul_101 = None
k_22[(slice(2, -2, None), slice(1, -2, None), slice(None, None, None))] = truediv_23; setitem_28 = k_22; k_22 = truediv_23 = None
zeros_like_9 = torch.zeros_like(k_11)
getitem_150 = zeros_like_9[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]; zeros_like_9 = None
zeros_like_10 = torch.zeros_like(k_11); k_11 = None
getitem_151 = zeros_like_10[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]; zeros_like_10 = None
getitem_152 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_153 = zeros_like_4[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_102 = getitem_152 * getitem_153; getitem_152 = getitem_153 = None
getitem_154 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_155 = zeros_like_5[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_103 = getitem_154 * getitem_155; getitem_154 = getitem_155 = None
add_41 = mul_102 + mul_103; mul_102 = mul_103 = None
getitem_156 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_157 = zeros_like[(slice(1, -3, None), slice(2, -2, None), slice(0, 25, None))]
mul_104 = getitem_156 * getitem_157; getitem_156 = getitem_157 = None
getitem_158 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_159 = zeros_like_1[(slice(1, -3, None), slice(2, -2, None), slice(0, 25, None))]
mul_105 = getitem_158 * getitem_159; getitem_158 = getitem_159 = None
add_42 = mul_104 + mul_105; mul_104 = mul_105 = None
neg_18 = -add_42; add_42 = None
tensor_16 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_9 = torch.min(tensor_16, add_41); tensor_16 = None
sub_17 = min_9 - 1e-20; min_9 = None
truediv_24 = neg_18 / sub_17; neg_18 = sub_17 = None
abs_11 = torch.abs(truediv_24)
neg_19 = -abs_11; abs_11 = None
add_43 = neg_19 + 0.001; neg_19 = None
truediv_25 = add_43 / 0.001; add_43 = None
tanh_8 = torch.tanh(truediv_25); truediv_25 = None
add_44 = 1.0 + tanh_8; tanh_8 = None
mul_106 = 0.5 * add_44; add_44 = None
getitem_160 = dxu[(slice(1, -3, None), None, None)]
getitem_161 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_107 = getitem_160 * getitem_161; getitem_160 = getitem_161 = None
mul_108 = mul_107 * mul_106; mul_107 = None
pow_1 = truediv_24 ** 2
mul_109 = mul_108 * pow_1; mul_108 = pow_1 = None
getitem_162 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_110 = mul_109 * getitem_162; mul_109 = getitem_162 = None
getitem_150 += mul_110; iadd_8 = getitem_150; getitem_150 = mul_110 = None
mul_111 = mul_106 * truediv_24; mul_106 = truediv_24 = None
getitem_163 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_112 = mul_111 * getitem_163; mul_111 = getitem_163 = None
ai_bx[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 0, 0)] = mul_112; setitem_29 = ai_bx; mul_112 = None
getitem_164 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_165 = zeros_like[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
mul_113 = getitem_164 * getitem_165; getitem_164 = getitem_165 = None
getitem_166 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_167 = zeros_like_1[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
mul_114 = getitem_166 * getitem_167; getitem_166 = getitem_167 = None
add_45 = mul_113 + mul_114; mul_113 = mul_114 = None
neg_20 = -add_45; add_45 = None
tensor_17 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_10 = torch.min(tensor_17, add_41); tensor_17 = None
sub_18 = min_10 - 1e-20; min_10 = None
truediv_26 = neg_20 / sub_18; neg_20 = sub_18 = None
abs_12 = torch.abs(truediv_26)
neg_21 = -abs_12; abs_12 = None
add_46 = neg_21 + 0.001; neg_21 = None
truediv_27 = add_46 / 0.001; add_46 = None
tanh_9 = torch.tanh(truediv_27); truediv_27 = None
add_47 = 1.0 + tanh_9; tanh_9 = None
mul_115 = 0.5 * add_47; add_47 = None
getitem_168 = dxu[(slice(2, -2, None), None, None)]
getitem_169 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_116 = getitem_168 * getitem_169; getitem_168 = getitem_169 = None
mul_117 = mul_116 * mul_115; mul_116 = None
pow_2 = truediv_26 ** 2
mul_118 = mul_117 * pow_2; mul_117 = pow_2 = None
getitem_170 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_119 = mul_118 * getitem_170; mul_118 = getitem_170 = None
iadd_8 += mul_119; iadd_9 = iadd_8; iadd_8 = mul_119 = None
mul_120 = mul_115 * truediv_26; mul_115 = truediv_26 = None
getitem_171 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_121 = mul_120 * getitem_171; mul_120 = getitem_171 = None
ai_bx[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 1, 0)] = mul_121; setitem_30 = ai_bx; mul_121 = None
getitem_172 = cosu[slice(1, -3, None)]
getitem_173 = dyu[slice(1, -3, None)]
mul_122 = getitem_172 * getitem_173; getitem_172 = getitem_173 = None
getitem_174 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_175 = zeros_like_2[(slice(2, -2, None), slice(1, -3, None), slice(0, 25, None))]
mul_123 = getitem_174 * getitem_175; getitem_174 = getitem_175 = None
getitem_176 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_177 = zeros_like_3[(slice(2, -2, None), slice(1, -3, None), slice(0, 25, None))]
mul_124 = getitem_176 * getitem_177; getitem_176 = getitem_177 = None
add_48 = mul_123 + mul_124; mul_123 = mul_124 = None
neg_22 = -add_48; add_48 = None
tensor_18 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_11 = torch.min(tensor_18, add_41); tensor_18 = None
sub_19 = min_11 - 1e-20; min_11 = None
truediv_28 = neg_22 / sub_19; neg_22 = sub_19 = None
abs_13 = torch.abs(truediv_28)
neg_23 = -abs_13; abs_13 = None
add_49 = neg_23 + 0.001; neg_23 = None
truediv_29 = add_49 / 0.001; add_49 = None
tanh_10 = torch.tanh(truediv_29); truediv_29 = None
add_50 = 1.0 + tanh_10; tanh_10 = None
mul_125 = 0.5 * add_50; add_50 = None
getitem_178 = mul_122[(None, slice(None, None, None), None)]; mul_122 = None
getitem_179 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_126 = getitem_178 * getitem_179; getitem_178 = getitem_179 = None
mul_127 = mul_126 * mul_125; mul_126 = None
pow_3 = truediv_28 ** 2
mul_128 = mul_127 * pow_3; mul_127 = pow_3 = None
getitem_180 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_129 = mul_128 * getitem_180; mul_128 = getitem_180 = None
getitem_151 += mul_129; iadd_10 = getitem_151; getitem_151 = mul_129 = None
mul_130 = mul_125 * truediv_28; mul_125 = truediv_28 = None
getitem_181 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_131 = mul_130 * getitem_181; mul_130 = getitem_181 = None
ai_by[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 0, 0)] = mul_131; setitem_31 = ai_by; mul_131 = None
getitem_182 = cosu[slice(2, -2, None)]
getitem_183 = dyu[slice(2, -2, None)]
mul_132 = getitem_182 * getitem_183; getitem_182 = getitem_183 = None
getitem_184 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_185 = zeros_like_2[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
mul_133 = getitem_184 * getitem_185; getitem_184 = getitem_185 = None
getitem_186 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
getitem_187 = zeros_like_3[(slice(2, -2, None), slice(2, -2, None), slice(0, 25, None))]
mul_134 = getitem_186 * getitem_187; getitem_186 = getitem_187 = None
add_51 = mul_133 + mul_134; mul_133 = mul_134 = None
neg_24 = -add_51; add_51 = None
tensor_19 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_12 = torch.min(tensor_19, add_41); tensor_19 = add_41 = None
sub_20 = min_12 - 1e-20; min_12 = None
truediv_30 = neg_24 / sub_20; neg_24 = sub_20 = None
abs_14 = torch.abs(truediv_30)
neg_25 = -abs_14; abs_14 = None
add_52 = neg_25 + 0.001; neg_25 = None
truediv_31 = add_52 / 0.001; add_52 = None
tanh_11 = torch.tanh(truediv_31); truediv_31 = None
add_53 = 1.0 + tanh_11; tanh_11 = None
mul_135 = 0.5 * add_53; add_53 = None
getitem_188 = mul_132[(None, slice(None, None, None), None)]; mul_132 = None
getitem_189 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_136 = getitem_188 * getitem_189; getitem_188 = getitem_189 = None
mul_137 = mul_136 * mul_135; mul_136 = None
pow_4 = truediv_30 ** 2
mul_138 = mul_137 * pow_4; mul_137 = pow_4 = None
getitem_190 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_139 = mul_138 * getitem_190; mul_138 = getitem_190 = None
iadd_10 += mul_139; iadd_11 = iadd_10; iadd_10 = mul_139 = None
mul_140 = mul_135 * truediv_30; mul_135 = truediv_30 = None
getitem_191 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_141 = mul_140 * getitem_191; mul_140 = getitem_191 = None
ai_by[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 1, 0)] = mul_141; setitem_32 = ai_by; mul_141 = None
getitem_192 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_193 = zeros_like_4[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]; zeros_like_4 = None
mul_142 = getitem_192 * getitem_193; getitem_192 = getitem_193 = None
getitem_194 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_195 = zeros_like_5[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]; zeros_like_5 = None
mul_143 = getitem_194 * getitem_195; getitem_194 = getitem_195 = None
add_54 = mul_142 + mul_143; mul_142 = mul_143 = None
getitem_196 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_197 = zeros_like[(slice(1, -3, None), slice(2, -2, None), slice(1, 26, None))]
mul_144 = getitem_196 * getitem_197; getitem_196 = getitem_197 = None
getitem_198 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_199 = zeros_like_1[(slice(1, -3, None), slice(2, -2, None), slice(1, 26, None))]
mul_145 = getitem_198 * getitem_199; getitem_198 = getitem_199 = None
add_55 = mul_144 + mul_145; mul_144 = mul_145 = None
neg_26 = -add_55; add_55 = None
tensor_20 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_13 = torch.min(tensor_20, add_54); tensor_20 = None
sub_21 = min_13 - 1e-20; min_13 = None
truediv_32 = neg_26 / sub_21; neg_26 = sub_21 = None
abs_15 = torch.abs(truediv_32)
neg_27 = -abs_15; abs_15 = None
add_56 = neg_27 + 0.001; neg_27 = None
truediv_33 = add_56 / 0.001; add_56 = None
tanh_12 = torch.tanh(truediv_33); truediv_33 = None
add_57 = 1.0 + tanh_12; tanh_12 = None
mul_146 = 0.5 * add_57; add_57 = None
getitem_200 = dxu[(slice(1, -3, None), None, None)]
getitem_201 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_147 = getitem_200 * getitem_201; getitem_200 = getitem_201 = None
mul_148 = mul_147 * mul_146; mul_147 = None
pow_5 = truediv_32 ** 2
mul_149 = mul_148 * pow_5; mul_148 = pow_5 = None
getitem_202 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_150 = mul_149 * getitem_202; mul_149 = getitem_202 = None
iadd_9 += mul_150; iadd_12 = iadd_9; iadd_9 = mul_150 = None
mul_151 = mul_146 * truediv_32; mul_146 = truediv_32 = None
getitem_203 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_152 = mul_151 * getitem_203; mul_151 = getitem_203 = None
ai_bx[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 0, 1)] = mul_152; setitem_33 = ai_bx; mul_152 = None
getitem_204 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_205 = zeros_like[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]; zeros_like = None
mul_153 = getitem_204 * getitem_205; getitem_204 = getitem_205 = None
getitem_206 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_207 = zeros_like_1[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]; zeros_like_1 = None
mul_154 = getitem_206 * getitem_207; getitem_206 = getitem_207 = None
add_58 = mul_153 + mul_154; mul_153 = mul_154 = None
neg_28 = -add_58; add_58 = None
tensor_21 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_14 = torch.min(tensor_21, add_54); tensor_21 = None
sub_22 = min_14 - 1e-20; min_14 = None
truediv_34 = neg_28 / sub_22; neg_28 = sub_22 = None
abs_16 = torch.abs(truediv_34)
neg_29 = -abs_16; abs_16 = None
add_59 = neg_29 + 0.001; neg_29 = None
truediv_35 = add_59 / 0.001; add_59 = None
tanh_13 = torch.tanh(truediv_35); truediv_35 = None
add_60 = 1.0 + tanh_13; tanh_13 = None
mul_155 = 0.5 * add_60; add_60 = None
getitem_208 = dxu[(slice(2, -2, None), None, None)]; dxu = None
getitem_209 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_156 = getitem_208 * getitem_209; getitem_208 = getitem_209 = None
mul_157 = mul_156 * mul_155; mul_156 = None
pow_6 = truediv_34 ** 2
mul_158 = mul_157 * pow_6; mul_157 = pow_6 = None
getitem_210 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_159 = mul_158 * getitem_210; mul_158 = getitem_210 = None
iadd_12 += mul_159; iadd_13 = iadd_12; iadd_12 = mul_159 = None
mul_160 = mul_155 * truediv_34; mul_155 = truediv_34 = None
getitem_211 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_161 = mul_160 * getitem_211; mul_160 = getitem_211 = None
ai_bx[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 1, 1)] = mul_161; setitem_34 = ai_bx; ai_bx = mul_161 = None
getitem_212 = cosu[slice(1, -3, None)]
getitem_213 = dyu[slice(1, -3, None)]
mul_162 = getitem_212 * getitem_213; getitem_212 = getitem_213 = None
getitem_214 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_215 = zeros_like_2[(slice(2, -2, None), slice(1, -3, None), slice(1, 26, None))]
mul_163 = getitem_214 * getitem_215; getitem_214 = getitem_215 = None
getitem_216 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]
getitem_217 = zeros_like_3[(slice(2, -2, None), slice(1, -3, None), slice(1, 26, None))]
mul_164 = getitem_216 * getitem_217; getitem_216 = getitem_217 = None
add_61 = mul_163 + mul_164; mul_163 = mul_164 = None
neg_30 = -add_61; add_61 = None
tensor_22 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_15 = torch.min(tensor_22, add_54); tensor_22 = None
sub_23 = min_15 - 1e-20; min_15 = None
truediv_36 = neg_30 / sub_23; neg_30 = sub_23 = None
abs_17 = torch.abs(truediv_36)
neg_31 = -abs_17; abs_17 = None
add_62 = neg_31 + 0.001; neg_31 = None
truediv_37 = add_62 / 0.001; add_62 = None
tanh_14 = torch.tanh(truediv_37); truediv_37 = None
add_63 = 1.0 + tanh_14; tanh_14 = None
mul_165 = 0.5 * add_63; add_63 = None
getitem_218 = mul_162[(None, slice(None, None, None), None)]; mul_162 = None
getitem_219 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_166 = getitem_218 * getitem_219; getitem_218 = getitem_219 = None
mul_167 = mul_166 * mul_165; mul_166 = None
pow_7 = truediv_36 ** 2
mul_168 = mul_167 * pow_7; mul_167 = pow_7 = None
getitem_220 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_169 = mul_168 * getitem_220; mul_168 = getitem_220 = None
iadd_11 += mul_169; iadd_14 = iadd_11; iadd_11 = mul_169 = None
mul_170 = mul_165 * truediv_36; mul_165 = truediv_36 = None
getitem_221 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_171 = mul_170 * getitem_221; mul_170 = getitem_221 = None
ai_by[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 0, 1)] = mul_171; setitem_35 = ai_by; mul_171 = None
getitem_222 = cosu[slice(2, -2, None)]; cosu = None
getitem_223 = dyu[slice(2, -2, None)]; dyu = None
mul_172 = getitem_222 * getitem_223; getitem_222 = getitem_223 = None
getitem_224 = mul_5[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]; mul_5 = None
getitem_225 = zeros_like_2[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]; zeros_like_2 = None
mul_173 = getitem_224 * getitem_225; getitem_224 = getitem_225 = None
getitem_226 = mul_7[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]; mul_7 = None
getitem_227 = zeros_like_3[(slice(2, -2, None), slice(2, -2, None), slice(1, 26, None))]; zeros_like_3 = None
mul_174 = getitem_226 * getitem_227; getitem_226 = getitem_227 = None
add_64 = mul_173 + mul_174; mul_173 = mul_174 = None
neg_32 = -add_64; add_64 = None
tensor_23 = torch.tensor([0.0], device = device(type='cuda', index=0))
min_16 = torch.min(tensor_23, add_54); tensor_23 = add_54 = None
sub_24 = min_16 - 1e-20; min_16 = None
truediv_38 = neg_32 / sub_24; neg_32 = sub_24 = None
abs_18 = torch.abs(truediv_38)
neg_33 = -abs_18; abs_18 = None
add_65 = neg_33 + 0.001; neg_33 = None
truediv_39 = add_65 / 0.001; add_65 = None
tanh_15 = torch.tanh(truediv_39); truediv_39 = None
add_66 = 1.0 + tanh_15; tanh_15 = None
mul_175 = 0.5 * add_66; add_66 = None
getitem_228 = mul_172[(None, slice(None, None, None), None)]; mul_172 = None
getitem_229 = k_iso[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]; k_iso = None
mul_176 = getitem_228 * getitem_229; getitem_228 = getitem_229 = None
mul_177 = mul_176 * mul_175; mul_176 = None
pow_8 = truediv_38 ** 2
mul_178 = mul_177 * pow_8; mul_177 = pow_8 = None
getitem_230 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]
mul_179 = mul_178 * getitem_230; mul_178 = getitem_230 = None
iadd_14 += mul_179; iadd_15 = iadd_14; iadd_14 = mul_179 = None
mul_180 = mul_175 * truediv_38; mul_175 = truediv_38 = None
getitem_231 = mask_w[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))]; mask_w = None
mul_181 = mul_180 * getitem_231; mul_180 = getitem_231 = None
ai_by[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None), 1, 1)] = mul_181; setitem_36 = ai_by; ai_by = mul_181 = None
getitem_232 = dxt[(slice(2, -2, None), None, None)]; dxt = None
mul_182 = 4 * getitem_232; getitem_232 = None
truediv_40 = iadd_13 / mul_182; iadd_13 = mul_182 = None
getitem_233 = dyt[(None, slice(2, -2, None), None)]; dyt = None
mul_183 = 4 * getitem_233; getitem_233 = None
getitem_234 = cost[(None, slice(2, -2, None), None)]; cost = None
mul_184 = mul_183 * getitem_234; mul_183 = getitem_234 = None
truediv_41 = iadd_15 / mul_184; iadd_15 = mul_184 = None
add_67 = truediv_40 + truediv_41; truediv_40 = truediv_41 = None
k_33[(slice(2, -2, None), slice(2, -2, None), slice(None, -1, None))] = add_67; setitem_37 = k_33; add_67 = None
k_33[(slice(2, -2, None), slice(2, -2, None), -1)] = 0.0; setitem_38 = k_33; k_33 = None
return ()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment