Test crashes sometimes for `scaled_dot_product_attention_bool_mask`

[justinchuby]

E.g. https://github.com/microsoft/onnx-script/actions/runs/4539041271/jobs/7998517646?pr=566#step:5:3923

worker 'gw1' crashed while running 'onnxscript/tests/function_libs/torch_lib/ops_correctness_test.py::TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__nn_functional_scaled_dot_product_attention_bool_mask_cpu_float32'

[justinchuby]

Profit!

TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__nn_functional_scaled_dot_product_attention_cpu_float32 - AssertionError: ONNX Runtime aborted:
Inputs:
{'input_0': array([[[ 3.6824992 ,  0.50817776, -7.971771  , -3.3470182 ,
          5.651129  ,  5.1625967 ,  3.1133838 ,  7.9440765 ],
        [-4.679297  , -2.6819472 , -1.34092   , -2.211886  ,
          7.7628407 , -8.993065  ,  1.7052691 ,  8.833969  ],
        [-5.00294   ,  7.152945  , -7.754232  ,  1.5221719 ,
         -5.1234965 ,  2.8475125 ,  1.9688144 ,  1.9786946 ]],

       [[-2.0416298 , -2.787754  , -8.715644  ,  1.4930023 ,
          8.591377  , -6.7724156 , -3.462263  ,  2.3617163 ],
        [ 4.6034536 ,  3.2687244 ,  5.831397  ,  1.1378338 ,
         -2.2142708 , -7.502145  , -8.995697  , -6.89208   ],
        [ 4.521224  ,  8.088249  ,  8.566948  ,  6.740285  ,
          0.4644953 ,  6.197366  ,  3.4216266 , -6.0938964 ]],

       [[-1.6890653 , -8.352889  ,  7.725936  ,  1.9562176 ,
          8.264622  , -1.0898427 ,  7.146616  ,  4.2126226 ],
        [-1.0984998 , -0.29764366,  6.4143376 , -1.1836277 ,
          6.9980903 ,  1.974403  , -3.525754  , -4.171672  ],
        [-4.377688  ,  5.7745514 ,  6.803203  , -4.2311287 ,
         -8.21504   ,  5.101013  , -2.0465417 , -8.263691  ]],

       [[ 1.8895293 , -0.8955177 ,  6.484178  ,  1.9106255 ,
         -5.548896  , -1.6276009 , -6.462818  ,  4.208257  ],
        [ 2.6327426 ,  5.3127975 , -6.6609645 , -5.009738  ,
         -5.4769325 , -7.4629097 , -8.394707  , -7.2591786 ],
        [ 6.269464  , -2.715213  , -1.6603786 ,  2.7127583 ,
          2.3688498 ,  4.0463376 , -6.3547025 ,  7.9970865 ]]],
      dtype=float32),
 'input_1': array([[[[-7.1899724e-01, -4.2738266e+00,  5.8091240e+00,
          -7.4829779e+00, -5.3650808e+00,  5.0721526e+00,
           3.4888351e+00, -3.5642920e+00],
         [-7.1924722e-01, -1.6436620e+00,  3.6597099e+00,
           1.8056395e+00,  3.2981794e+00,  7.4339695e+00,
          -6.1898794e+00, -6.9231420e+00],
         [-2.1351576e+00, -4.0983467e+00, -6.4004545e+00,
           8.3525362e+00,  7.5186548e+00,  7.2669821e+00,
          -5.5431499e+00, -1.0642394e+00],
         [-7.0026522e+00, -3.2450545e+00,  7.4999037e+00,
           1.1283410e+00,  3.1221836e+00, -8.9466047e+00,
           3.7078857e+00,  6.0683627e+00],
         [-3.4560103e+00,  6.8391814e+00,  4.0536451e+00,
          -6.5714669e+00, -6.7465944e+00,  1.2635243e-01,
          -8.6807718e+00, -6.2895775e+00],
         [ 3.5202608e+00, -2.9611964e+00, -3.1890929e-01,
          -3.6772399e+00,  5.8250852e+00,  7.3904290e+00,
           6.4942417e+00, -8.0034370e+00]],

        [[-2.6608105e+00,  4.3589478e+00,  5.3230801e+00,
          -2.0171092e+00,  4.9280987e+00, -2.6609015e-01,
           9.9968612e-01,  6.0641217e+00],
         [ 6.1567974e+00, -6.3537207e+00, -5.4640074e+00,
           3.1882658e+00,  5.4091339e+00, -2.0955768e+00,
           7.9811831e+00, -1.0368057e+00],
         [-4.5640469e+00, -9.1866410e-01, -7.6602378e+00,
           2.7885296e+00, -8.2743692e+00,  8.3511810e+00,
          -6.8289633e+00,  5.8585949e+00],
         [-6.7507863e+00, -1.3116463e+00, -4.9842844e+00,
           4.5988898e+00,  5.2198944e+00, -8.4159822e+00,
          -7.7264547e+00,  1.1848015e+00],
         [-1.9361150e-01,  5.2183986e-02,  2.7360742e+00,
          -1.7787284e+00,  7.5948143e+00,  6.4017034e-01,
          -4.1461620e+00,  3.1228895e+00],
         [-6.9828029e+00, -4.3042555e+00,  2.6239922e+00,
          -3.7758765e+00, -8.5922298e+00,  4.5044374e+00,
           2.3163772e+00,  2.6195087e+00]],

        [[ 3.1282616e+00, -8.8809557e+00, -3.9208927e+00,
          -4.3719063e+00,  8.6117878e+00, -3.7917004e+00,
          -5.9125972e+00, -3.1587839e+00],
         [ 8.6476243e-01,  3.6269636e+00,  6.6279058e+00,
           6.4973092e+00,  9.2429030e-01,  8.1990528e+00,
           4.7852411e+00,  4.8019214e+00],
         [-7.7175350e+00,  7.3527718e-01, -1.9114966e+00,
          -8.5530787e+00, -4.4446745e+00,  6.8478780e+00,
           3.3004723e+00,  8.3176727e+00],
         [ 6.1366301e+00, -6.8849087e-02,  8.0888424e+00,
           1.8983184e+00, -1.3220093e+00, -3.8486180e+00,
           5.8009434e+00,  4.0348072e+00],
         [ 6.2488647e+00, -1.5634747e+00, -8.3295412e+00,
           1.6320630e+00,  9.8490608e-01, -5.8582330e+00,
           2.1819921e+00, -1.0892472e+00],
         [ 7.6784315e+00,  4.2354012e-01, -4.4449911e+00,
           7.2103410e+00,  2.6743126e-01, -1.5115514e+00,
           4.2986393e+00,  8.6815367e+00]],

        [[ 4.1037550e+00,  2.5377774e-01,  7.3711600e+00,
          -3.9469852e+00,  7.9539547e+00, -5.1278191e+00,
           1.6490865e-01,  8.1909714e+00],
         [ 3.5840235e+00, -1.8861337e+00,  5.5448837e+00,
           5.3087587e+00, -8.1125641e+00, -1.0030915e+00,
          -7.0829158e+00, -1.9399216e+00],
         [-1.0848806e+00, -3.6003990e+00, -3.2893891e+00,
           4.8675194e+00, -2.6728535e-01, -8.9547920e+00,
           1.5470392e+00, -8.8509493e+00],
         [-6.1742187e+00,  5.5013180e+00,  6.8800001e+00,
          -3.1932116e+00, -5.4649558e+00,  4.0727711e+00,
          -8.5974188e+00,  5.5524559e+00],
         [ 6.1993542e+00,  3.7977183e+00,  2.7468739e+00,
           1.8610615e+00, -6.8995171e+00, -3.6050830e+00,
           5.8053646e+00,  7.5802135e+00],
         [ 2.4039245e-01,  2.3122122e+00, -8.5365381e+00,
           6.3534451e+00, -5.2420683e+00, -4.8244820e+00,
           8.8806286e+00, -2.5799255e+00]]],


       [[[-4.5469117e+00,  1.2379907e+00,  7.0926232e+00,
           8.1777992e+00,  7.5339389e-01,  6.1271234e+00,
           6.4185357e+00, -1.5041968e+00],
         [-7.1115713e+00,  2.4213588e-01, -3.0572934e+00,
           3.1539569e+00, -6.5942035e+00, -6.0368500e+00,
          -6.7779636e+00, -6.7329454e+00],
         [ 6.8458652e+00, -8.9680138e+00, -3.3129036e-01,
          -5.5545340e+00, -3.7854156e+00, -2.5936584e+00,
           4.2222133e+00,  6.0517912e+00],
         [-2.0249572e+00,  2.9347765e-01, -9.8998189e-02,
           8.1959581e+00,  4.0224543e+00, -1.8686993e+00,
           6.0943651e+00, -7.2670989e+00],
         [ 3.5670311e+00,  6.9814262e+00, -4.6712914e+00,
           3.3676667e+00,  4.1266775e+00, -7.5507317e+00,
           2.5889549e+00,  6.8061590e+00],
         [ 8.2649374e-01, -1.7394587e+00, -3.5270941e+00,
          -5.9733334e+00,  2.1856475e+00, -8.5077410e+00,
          -7.7038922e+00, -3.2564583e+00]],

        [[-8.5227022e+00,  8.4039927e+00, -2.8999648e+00,
          -4.7849803e+00,  2.2623906e+00,  3.8732815e+00,
           6.7776151e+00, -3.6334953e+00],
         [-1.5726467e+00, -6.7962456e+00, -7.8587823e+00,
          -3.9159942e+00,  1.0126208e+00, -4.4616885e+00,
          -2.4844542e+00,  3.2988973e+00],
         [ 1.8569909e+00,  5.5531983e+00,  4.1560197e+00,
           4.7458696e+00, -4.6527410e+00,  4.8272972e+00,
          -8.0666418e+00,  2.7715511e+00],
         [-5.6761465e+00,  2.4851997e+00, -8.2341690e+00,
          -1.4906259e+00,  5.5357599e+00, -4.5007381e+00,
           5.5714383e+00, -1.3993331e+00],
         [ 8.4433994e+00,  8.4793863e+00, -5.3101525e+00,
          -2.1079812e+00, -1.8394611e+00, -3.4029670e+00,
           2.9584830e+00, -8.1945066e+00],
         [-5.3972964e+00, -1.4108773e+00, -8.0075903e+00,
          -8.6204233e+00,  3.6909633e+00, -7.5145659e+00,
           5.9796000e+00, -5.2258339e+00]],

        [[-7.1778381e-01,  5.4267430e+00, -8.8753090e+00,
          -8.7624407e+00,  3.7477648e-01, -3.6541085e+00,
           4.2068701e+00,  3.1413722e+00],
         [ 1.1077191e+00,  3.3783140e+00,  7.6997800e+00,
           1.4284147e+00,  5.9547205e+00, -5.0743222e+00,
           7.4154592e+00, -1.6003089e+00],
         [ 4.8724899e+00, -7.4795189e+00, -5.4084940e+00,
          -5.2685289e+00,  4.7543840e+00, -2.9990745e-01,
          -8.2597904e+00,  8.8265047e+00],
         [ 2.3457463e+00,  8.2777042e+00,  1.1234658e+00,
           1.2901769e+00,  6.4872026e+00, -6.1230860e+00,
           8.4994974e+00, -5.7886233e+00],
         [ 8.1574559e-01, -5.8533630e+00, -5.9631739e+00,
           6.9272938e+00, -2.3770370e+00, -1.6890750e+00,
           7.0318232e+00,  8.8172703e+00],
         [-8.1174850e+00, -1.5709591e-01,  2.5699918e+00,
           8.6411610e+00, -3.5880361e+00, -2.9290185e+00,
          -5.2750683e+00,  2.2996025e+00]],

        [[ 2.2196717e+00, -1.4192212e-01, -9.1280508e-01,
          -6.2656488e+00, -3.6966140e+00, -3.4228883e+00,
           2.3640594e+00,  3.9744167e+00],
         [ 4.2016773e+00, -6.8698926e+00, -6.9375014e-01,
          -1.7175858e+00,  8.9774923e+00,  6.4945574e+00,
          -6.7181258e+00, -4.9452114e+00],
         [-6.7987399e+00, -2.2524323e+00, -7.8746104e+00,
           4.2298632e+00, -3.9691062e+00, -4.5169935e+00,
           4.2535114e-01, -2.8756618e+00],
         [ 6.7990189e+00,  8.1583271e+00, -8.0357952e+00,
           4.0352116e+00,  1.8779433e-01,  4.7238102e+00,
          -3.3740537e+00, -6.8142700e-01],
         [-1.6835368e-01,  3.0858943e+00,  4.4873691e+00,
           3.5100384e+00,  1.1045508e+00, -7.1797361e+00,
          -5.1700535e+00, -5.1272292e+00],
         [-8.9690094e+00,  8.6211004e+00, -4.0746884e+00,
          -5.9822435e+00, -4.0757174e+00,  4.8867378e+00,
          -5.6418905e+00, -7.5201297e+00]]],


       [[[-6.7268429e+00,  1.9778920e+00, -3.6821258e-01,
          -8.1202641e+00, -6.1563187e+00, -7.7996106e+00,
           3.0874865e+00,  4.0729170e+00],
         [ 1.5330713e+00, -2.8834391e+00, -4.8220387e+00,
           3.5580716e+00, -2.0667257e+00, -1.4447011e+00,
          -2.9503708e+00,  1.5641710e+00],
         [ 2.1111710e+00,  7.3944802e+00,  7.9928842e+00,
          -2.3195195e+00,  1.3375951e+00, -6.0156112e+00,
          -6.7563653e+00,  7.4631319e+00],
         [ 1.0601882e+00, -4.4749684e+00,  4.5246806e+00,
           3.6308193e+00,  3.2216840e+00,  3.8081317e+00,
           4.8066580e-01, -6.5470457e+00],
         [-8.0008736e+00,  8.7865744e+00,  6.2830338e+00,
          -7.0215459e+00,  7.0528321e+00, -8.4174728e+00,
           5.9867773e+00,  6.2671132e+00],
         [ 5.6542921e+00,  3.8089805e+00, -5.4511852e+00,
          -3.8596795e+00,  1.9293430e+00,  3.5495539e+00,
           7.8894196e+00,  4.0034115e-01]],

        [[ 3.7061400e+00, -6.7969270e+00, -1.8267635e+00,
           5.3521914e+00,  7.3558080e-01,  7.1148372e+00,
          -7.2218165e+00, -6.0053234e+00],
         [ 2.8293808e+00, -3.7635362e-01, -7.6900253e+00,
           3.8401389e+00,  3.7094307e-01,  2.1764452e+00,
           1.9777976e+00, -6.2058625e+00],
         [-6.0370998e+00, -7.2009354e+00,  6.5371494e+00,
          -4.9404240e+00,  8.1020813e+00, -2.8741393e+00,
           1.9920863e+00,  8.5957775e+00],
         [-4.8297744e+00, -8.5363169e+00,  1.2978941e+00,
          -6.4928789e+00,  4.9591846e+00,  3.5780067e+00,
           6.5208597e+00, -9.3491828e-01],
         [-5.8799610e+00, -8.2938938e+00, -5.8130808e+00,
           8.0164347e+00, -1.9266254e+00, -1.6485178e-01,
          -1.9096481e+00, -4.5952625e+00],
         [ 7.3677340e+00, -2.3317699e+00, -4.3680472e+00,
           8.3155432e+00,  5.3867888e+00,  4.7965374e+00,
           6.4495573e+00, -5.7153873e+00]],

        [[ 2.9309044e+00, -2.5048571e+00,  3.2170415e-01,
           4.2393732e+00, -2.9874907e+00,  8.8308706e+00,
           4.4137359e-02,  5.1857948e+00],
         [ 8.2822859e-01, -8.6744375e+00, -5.3827901e+00,
          -7.4995098e+00,  1.0575017e+00, -6.0375986e+00,
          -7.1573448e+00,  1.9295082e+00],
         [-7.8780684e+00, -4.8974295e+00, -5.8221865e+00,
           4.3930702e+00,  8.2517281e+00,  3.1489263e+00,
          -5.0369902e+00,  6.8121619e+00],
         [ 8.3537655e+00, -2.4384992e+00,  6.6410818e+00,
           2.2752643e+00, -6.2602963e+00, -2.3506441e+00,
          -7.1778007e+00,  2.7155423e+00],
         [ 4.9769001e+00, -8.8521576e+00, -3.6783643e+00,
           1.9503844e-01,  7.2066655e+00, -4.9039664e+00,
          -5.2857006e-01, -2.5451589e-01],
         [-4.5259004e+00, -4.0339909e+00, -2.4080551e+00,
          -2.6892462e+00, -2.8968964e+00,  8.6046231e-01,
          -7.5238860e-01, -4.3800774e+00]],

        [[ 5.4502296e+00,  6.4385505e+00, -1.4108773e+00,
          -1.6989412e+00, -8.7716875e+00, -6.5796680e+00,
          -4.0990806e-01, -1.7158402e+00],
         [ 2.6331182e+00, -4.4771037e+00, -1.8640859e+00,
           1.3457490e+00,  8.8344049e+00,  7.0809565e+00,
          -5.9202008e+00, -7.4959481e-01],
         [-1.4734726e+00, -2.8769267e-01,  3.7493839e+00,
          -3.6205091e+00, -1.9184318e+00, -8.9154139e+00,
          -3.1971512e+00, -8.8269472e+00],
         [ 3.1860952e+00, -6.9278593e+00, -4.4472613e+00,
          -6.5265388e+00,  6.8438268e+00, -2.6420178e+00,
           7.8211799e+00, -4.4721718e+00],
         [ 6.8134761e+00, -1.7696207e+00, -6.7608910e+00,
           7.5262947e+00, -6.1228247e+00, -5.2924652e+00,
           4.9320297e+00, -8.4018159e-01],
         [-1.3669932e-01,  8.9240580e+00, -7.1422701e+00,
          -3.2507396e+00,  3.6613603e+00, -1.7330064e+00,
           1.5224218e-01, -3.0426936e+00]]],


       [[[-1.3942112e+00,  7.1839032e+00, -1.6729420e+00,
           3.6195810e+00,  6.7906194e+00,  6.9673986e+00,
           4.2673912e+00,  6.1863594e+00],
         [-2.3345013e+00,  7.1762502e-01, -3.8605795e+00,
          -8.7284765e+00,  7.5727806e+00, -8.6042204e+00,
          -5.6720405e+00, -3.2076559e+00],
         [-8.2325716e+00, -6.9842682e+00, -4.7076716e+00,
           1.0650001e+00, -7.8689532e+00, -6.4693575e+00,
          -5.2609997e+00, -6.7677979e+00],
         [ 2.1124778e+00,  3.0450153e-01,  1.9749856e+00,
           8.1192884e+00, -3.9406061e+00, -3.3742983e+00,
           2.8711147e+00,  7.8933926e+00],
         [ 2.8725107e+00, -7.7150846e-01,  2.9235039e+00,
          -8.8649769e+00, -6.1562929e+00, -3.7844102e+00,
          -8.1471882e+00, -6.7970886e+00],
         [-6.9216795e+00,  2.5836787e+00,  2.7266221e+00,
          -2.1676862e-01, -2.0801926e+00,  2.4266374e-01,
           5.4382429e+00, -6.7768016e+00]],

        [[ 6.3940697e+00,  3.9143558e+00, -3.1416392e+00,
           6.3145533e+00, -7.4543953e-03, -4.1244521e+00,
           8.1106386e+00,  2.6600432e+00],
         [ 6.7920012e+00, -2.4422779e+00, -7.7121687e+00,
           7.2174406e+00,  5.2065647e-01,  4.2741752e+00,
          -2.3327858e+00,  4.2567835e+00],
         [-7.1112967e+00,  3.1746733e-01,  6.8336978e+00,
          -2.4167197e+00,  7.7300525e+00,  3.3097420e+00,
           4.4366388e+00,  8.8930531e+00],
         [-2.1897941e+00, -2.0952549e+00,  5.6594954e+00,
          -4.5546107e+00, -6.6358676e+00,  5.8684425e+00,
          -7.8526454e+00, -4.2624216e+00],
         [ 2.1807680e+00,  1.3253288e+00, -5.4561248e+00,
          -2.2769079e+00,  5.7771726e+00, -5.8317447e+00,
           5.1623950e+00, -9.6378958e-01],
         [-3.2215466e+00, -8.5734272e+00, -2.3606281e+00,
           7.4004955e+00, -4.9304538e+00, -7.1427946e+00,
          -7.0218408e-01,  6.3800020e+00]],

        [[-2.4413552e+00,  8.8226805e+00,  8.6688747e+00,
           8.2456303e+00, -5.1477995e+00, -6.4066162e+00,
          -8.4346704e+00, -1.7500427e+00],
         [-5.8611622e+00,  4.1928344e+00,  8.1168687e-01,
           7.1380105e+00, -8.3490257e+00, -7.5629129e+00,
          -1.6748067e+00,  2.3720598e+00],
         [ 7.5774841e+00, -7.5245619e-01,  7.6912527e+00,
          -2.4364393e+00, -1.2704905e+00, -3.4654548e+00,
          -3.8250360e+00, -3.4094546e+00],
         [-7.7797914e+00, -1.1047976e+00, -3.0346942e+00,
          -5.8269124e+00, -6.6757545e+00, -8.6790047e+00,
          -2.0167615e+00, -7.3273220e+00],
         [ 8.0918541e+00, -2.4159055e+00,  2.4191928e+00,
          -6.4524441e+00,  9.1833794e-01, -8.4585638e+00,
           4.0517311e+00,  4.2503734e+00],
         [-8.7459154e+00,  3.9245460e+00,  1.1458172e+00,
          -5.6210508e+00, -4.1205220e+00,  8.8462982e+00,
          -4.6813412e+00,  7.8744593e+00]],

        [[ 1.8111606e+00,  4.1234994e+00, -4.5300570e+00,
           2.9255669e+00,  1.9378874e+00,  7.1779141e+00,
          -8.8999577e+00, -3.2721491e+00],
         [ 2.2411036e-01, -6.4597197e+00,  1.5338030e+00,
           5.8029237e+00,  5.4236789e+00, -4.6408010e+00,
           3.6380281e+00, -6.5486774e+00],
         [ 5.5946975e+00, -7.3559809e+00, -8.1593246e+00,
           3.1234422e+00, -8.4961116e-01,  2.2184992e+00,
           8.3744087e+00,  8.6216440e+00],
         [ 5.0746579e+00, -8.1641111e+00, -1.3856753e+00,
           9.8382568e-01, -8.7636156e+00,  7.5265608e+00,
           5.9059334e+00, -4.1574841e+00],
         [-1.8782684e+00, -5.2204351e+00,  3.4616139e+00,
           1.5521868e+00, -3.1202180e+00,  4.8116479e+00,
          -1.9454974e+00,  9.6651042e-01],
         [ 6.2313490e+00, -4.2285423e+00, -8.3006687e+00,
           4.8575330e+00, -5.8817658e+00,  3.4023521e+00,
          -2.7057729e+00,  6.0415225e+00]]]], dtype=float32),
 'input_2': array([[[[-0.34230995, -4.0978384 , -6.349059  ,  6.047771  ,
          -5.631379  ,  2.551255  ,  6.882798  , -7.2295017 ],
         [-7.553694  ,  0.91815555, -3.5329392 , -6.606769  ,
           5.9512577 , -8.372425  , -4.1034985 , -7.4758034 ],
         [ 2.3787193 ,  0.7762345 ,  3.754082  ,  3.768081  ,
          -0.0356015 , -2.702716  ,  7.9581776 ,  0.3742733 ],
         [ 2.1427407 , -7.4108696 , -8.820838  , -4.395982  ,
          -4.7160378 ,  2.6075492 , -4.000244  ,  0.93990827],
         [ 8.340205  , -2.6548986 ,  2.2993689 ,  0.11037505,
           5.922707  ,  3.936295  , -0.41812956,  4.922632  ],
         [-0.45790136, -1.5300007 , -3.5081182 , -7.342042  ,
           7.5179358 , -5.703294  , -8.579783  , -4.036794  ]],

        [[ 1.8180141 , -0.81675303,  4.0971346 ,  5.156386  ,
           5.2929974 , -7.5512433 ,  3.6692598 ,  4.4550023 ],
         [-4.9928417 ,  0.04359448,  1.3489387 ,  1.1285588 ,
           4.8295083 ,  7.6410804 ,  8.644264  ,  2.0150535 ],
         [ 5.263167  ,  3.8479776 , -4.8689647 ,  6.0192933 ,
          -0.63945794,  3.672884  , -5.973283  ,  7.200478  ],
         [ 2.008657  ,  2.0394938 , -8.0063925 , -0.78217936,
          -3.147153  ,  8.497046  , -8.23968   , -2.220986  ],
         [-4.4639134 , -6.524279  ,  3.4615722 , -6.213848  ,
           4.254297  ,  1.0802768 , -4.764764  , -1.024209  ],
         [ 0.7804874 , -4.010984  , -5.782582  , -5.071845  ,
           2.171483  , -7.930274  ,  5.6670957 ,  5.9107027 ]],

        [[ 8.275736  ,  7.6145887 , -5.0580616 , -3.905322  ,
           6.5873466 , -1.164304  , -6.626468  , -4.31118   ],
         [ 6.31756   , -5.1179714 ,  3.165544  , -7.5146894 ,
           5.9429803 ,  2.475356  ,  6.2080154 ,  0.306952  ],
         [ 4.568535  ,  7.375328  , -1.1353394 , -6.80179   ,
          -6.3504    ,  5.276816  , -5.7391677 , -7.878265  ],
         [-7.620439  , -3.9857883 ,  1.6696558 , -8.154997  ,
          -4.170337  ,  8.616638  , -1.5307474 , -3.037105  ],
         [ 4.2566166 , -7.4061036 , -7.0274777 ,  1.2868155 ,
           2.8450866 , -8.869081  ,  7.9074764 , -3.504482  ],
         [ 1.3749079 ,  6.397666  ,  3.4449348 ,  2.7945228 ,
          -8.211373  , -8.39686   ,  3.3470879 ,  0.3344618 ]],

        [[ 5.526359  ,  4.50389   , -7.2203007 ,  1.6565012 ,
          -7.0605087 , -4.975044  ,  2.4733067 , -0.503073  ],
         [-5.747767  ,  3.0517874 ,  6.24559   , -8.845017  ,
           8.441818  ,  4.968251  , -0.7585845 , -8.727442  ],
         [-3.8099234 ,  2.746667  , -3.834444  , -2.958364  ,
          -3.7167199 ,  7.738069  ,  7.6167045 , -7.0197434 ],
         [-2.7766004 ,  8.151881  , -0.5183884 , -4.3818927 ,
           1.0883235 ,  7.1210847 , -1.1453065 , -8.492519  ],
         [-3.0083947 ,  4.479186  , -1.3297963 ,  6.358906  ,
           3.0952597 ,  0.0789299 ,  3.437828  , -1.2148347 ],
         [-7.7152843 , -7.713045  , -7.0792937 , -1.4334733 ,
          -6.5075026 ,  1.3189505 ,  1.0319617 ,  6.8479414 ]]],


       [[[ 0.5060556 ,  1.5339446 ,  3.2002883 , -2.544455  ,
          -4.484561  , -8.466113  , -8.765525  ,  0.59179723],
         [-4.2559876 ,  2.4915285 ,  2.2367692 ,  7.0193615 ,
          -1.0987122 , -0.0380584 ,  6.211625  ,  1.5262123 ],
         [ 4.119171  , -4.8599524 , -6.5050607 ,  6.861494  ,
           6.415559  , -3.4077444 ,  1.4950504 , -8.004093  ],
         [ 0.629861  , -8.355907  ,  5.086895  , -8.920359  ,
          -2.1710293 ,  8.396976  , -3.0490515 , -5.216814  ],
         [ 8.11894   , -4.5523977 ,  1.9849752 ,  5.029893  ,
           2.418108  , -6.2179594 ,  8.168247  ,  4.3157053 ],
         [-5.65954   , -5.1365924 , -2.0757604 ,  3.0025828 ,
           3.4030323 ,  1.1588559 ,  5.9023085 ,  8.953258  ]],

        [[ 1.8515954 ,  3.525226  ,  5.118249  ,  1.7841175 ,
           0.5548954 ,  8.092021  ,  6.2133555 ,  5.8781395 ],
         [-3.5019822 , -4.3029294 , -4.880156  , -7.5367327 ,
          -0.91440153,  5.070908  ,  1.3185772 , -3.1836286 ],
         [-1.2450696 , -6.896912  , -5.781066  ,  0.82857835,
          -5.229164  ,  1.1092318 ,  7.422377  ,  2.30473   ],
         [ 6.889682  , -5.7019978 ,  0.23406887, -8.8410225 ,
           0.06235921, -6.630025  ,  8.446856  , -7.6939692 ],
         [ 3.07654   , -5.3836336 ,  2.491255  , -0.04694295,
          -1.151144  , -3.4929657 ,  5.7868576 , -1.62457   ],
         [ 0.767115  , -2.7122412 , -6.558548  , -3.7508407 ,
           8.069345  ,  8.56679   , -4.181401  ,  0.89693284]],

        [[ 1.4608608 ,  7.9395723 ,  1.2559626 ,  0.41896963,
          -2.6179614 ,  1.8980459 , -8.961205  , -8.979448  ],
         [-1.0432377 , -6.242316  , -6.250534  , -1.5276393 ,
           4.9933357 ,  4.4542627 , -3.1269214 , -0.6360065 ],
         [ 8.355031  ,  1.9468868 ,  6.120341  , -3.7278447 ,
          -8.053922  , -4.172763  ,  0.7909266 , -2.1797915 ],
         [-4.914157  ,  8.654715  , -8.307294  ,  0.22850919,
           3.586665  , -2.4280643 ,  3.9034038 ,  0.03194726],
         [-5.3498836 , -4.8937836 , -4.04868   , -8.787451  ,
           5.609341  ,  8.249806  , -4.623246  ,  0.4448465 ],
         [ 3.007388  ,  2.5482736 , -4.995444  , -5.4690695 ,
           8.761045  ,  5.240899  ,  8.413323  ,  6.3179026 ]],

        [[-0.6696371 ,  6.336029  , -0.28866792, -0.0558393 ,
          -3.8995607 ,  7.0187025 , -1.9134021 ,  1.3067003 ],
         [-7.176584  ,  3.0873265 , -4.2791715 , -2.0322926 ,
          -0.8047968 ,  3.3364782 ,  0.3673811 ,  1.56819   ],
         [-5.8456073 , -6.80033   ,  3.2412179 , -3.706018  ,
          -7.759484  , -3.714308  , -8.829413  ,  1.7982377 ],
         [ 0.6684011 ,  8.311016  , -4.962897  , -4.6545744 ,
           2.1489978 , -2.3648608 ,  7.0137653 ,  1.3691282 ],
         [-0.19555664, -3.4528592 ,  5.4378805 , -4.212881  ,
           3.6190457 ,  5.7816677 , -7.059673  , -7.9743114 ],
         [-5.075494  ,  4.318495  , -4.964339  ,  0.8507484 ,
           1.3588833 ,  1.8894832 , -4.589383  ,  0.3146907 ]]],


       [[[ 6.2364674 ,  4.665291  ,  1.7193228 , -7.157478  ,
           4.2342606 , -3.583723  ,  6.2986712 ,  3.7378986 ],
         [-1.6520036 ,  3.4791384 , -3.7425292 ,  1.5534518 ,
          -0.44907153,  0.14281583,  3.1689956 ,  8.764865  ],
         [-3.9725296 ,  8.788686  , -1.6515831 ,  8.559981  ,
           2.0253005 ,  0.7339339 , -3.0099182 , -4.742607  ],
         [-8.850726  , -4.9126387 , -0.80580103,  7.571518  ,
           4.8703218 ,  1.5691963 , -8.71032   ,  6.340592  ],
         [ 6.196273  , -4.0458193 ,  3.6451972 , -7.4987164 ,
          -4.5503087 , -2.716004  , -8.160566  , -0.5498121 ],
         [ 1.8430274 ,  2.6196856 ,  8.422474  ,  5.1977305 ,
          -8.234614  , -2.8249798 ,  5.042182  , -8.638262  ]],

        [[-0.5566045 ,  3.4055567 , -6.0347867 ,  7.7202396 ,
          -5.933469  , -6.6021404 , -6.831831  , -1.5764834 ],
         [ 3.0054526 ,  2.5924494 , -4.446918  ,  7.871669  ,
          -4.823057  ,  7.061883  ,  3.023122  , -2.904017  ],
         [ 6.656292  ,  2.9841828 , -7.652502  ,  4.5741286 ,
           8.511661  , -6.7322598 , -1.897414  , -7.2625446 ],
         [ 6.033968  ,  4.775005  ,  6.5643225 ,  2.5277867 ,
          -1.4840738 , -0.8781692 , -6.7057233 ,  5.7115536 ],
         [-3.6106932 ,  0.9877299 ,  0.30608404,  2.7207212 ,
          -7.229786  , -1.4255319 ,  4.15208   , -6.652448  ],
         [-3.866281  , -4.9557824 , -7.6150084 ,  8.787146  ,
           7.7758656 ,  5.639732  , -7.460018  , -0.08213997]],

        [[ 1.6527236 ,  3.9922428 ,  1.8127184 , -5.516509  ,
           0.2853613 , -7.2778115 ,  4.5133758 ,  5.669015  ],
         [-3.2203126 , -7.8779235 ,  0.08367848,  4.6010675 ,
          -6.14085   ,  3.799644  ,  6.268093  , -7.460702  ],
         [ 4.977277  ,  4.0702305 , -4.6315937 ,  7.553576  ,
          -4.3281403 ,  6.840512  , -6.898634  ,  7.098889  ],
         [-2.0235925 , -7.9501467 ,  4.7939253 , -0.73716974,
           2.5420046 ,  4.1460257 ,  0.9609475 ,  2.8745513 ],
         [-7.585773  , -7.9436517 , -4.2263126 , -4.7055035 ,
          -8.483265  ,  1.4052114 , -0.7140695 , -1.9820827 ],
         [-0.4182197 ,  4.3249044 , -3.3396592 ,  0.56857896,
          -0.8039224 ,  4.3460445 , -0.9574542 ,  3.2329597 ]],

        [[-3.523854  ,  6.971783  ,  1.5148151 ,  1.5140405 ,
           4.357374  , -8.971563  ,  3.9768887 ,  7.230226  ],
         [ 5.390065  ,  5.844248  , -7.2293396 ,  7.7898884 ,
           2.0122254 , -8.772329  ,  8.143966  ,  7.018045  ],
         [-1.4358433 ,  8.825392  ,  0.5666628 ,  7.2184963 ,
          -7.9558535 , -8.390788  ,  2.5075521 , -0.06732452],
         [-7.6399775 , -5.0466223 ,  5.0448103 ,  5.7374845 ,
          -6.8261423 , -4.0237017 ,  0.05233955, -0.01148844],
         [-6.628728  ,  6.655866  , -6.2843213 ,  6.4027185 ,
          -5.566676  , -8.942067  ,  2.2423835 , -4.305365  ],
         [ 6.067338  , -2.9838877 ,  0.8593594 ,  6.6992774 ,
          -5.8451815 , -3.7053301 ,  6.127894  ,  7.5919476 ]]],


       [[[ 1.5686717 ,  6.9187164 ,  5.444451  ,  5.0889187 ,
          -0.35480368,  8.689396  , -2.861092  ,  0.612609  ],
         [-2.7166636 , -0.6744189 ,  5.4085417 ,  6.823352  ,
           1.3416077 ,  5.916538  ,  6.3087487 ,  4.37452   ],
         [-2.4527354 ,  5.251175  ,  7.5547123 ,  7.4452276 ,
           2.331995  ,  6.6418686 , -6.217531  ,  1.7046597 ],
         [-7.907222  ,  4.620072  ,  0.41317284, -1.731353  ,
           5.737301  , -1.3776695 ,  3.6747088 ,  4.8420577 ],
         [-4.46467   , -7.796344  ,  0.8038162 , -5.5036    ,
          -5.1038737 ,  4.238401  ,  2.138669  , -2.0001135 ],
         [-3.271941  , -6.6557193 , -3.8459122 ,  0.57687986,
          -5.2578635 , -8.170835  , -8.719526  , -8.48733   ]],

        [[-4.1398983 ,  2.4248843 , -4.115657  ,  5.945626  ,
           4.9255066 , -7.0417085 ,  2.215833  , -3.8760872 ],
         [ 7.4642477 ,  6.81977   ,  4.1280193 , -4.5070915 ,
           1.3013864 , -5.674239  ,  1.3475128 ,  8.897213  ],
         [-1.7353356 , -5.1756988 , -4.387149  , -5.5859427 ,
           4.559204  , -5.3083544 ,  6.2997246 ,  5.254262  ],
         [ 5.4144382 , -0.43160605, -1.8787909 ,  3.2602606 ,
          -1.3443532 ,  3.6168933 ,  2.0089831 ,  3.1299973 ],
         [-3.8472211 ,  5.579997  , -1.7262236 ,  5.113995  ,
          -0.5548074 , -8.845102  , -0.7667781 , -2.8695989 ],
         [-4.4622483 , -0.3521011 , -3.4386177 ,  0.70515597,
          -2.6571198 , -5.584134  ,  3.3040287 , -7.370762  ]],

        [[-6.967508  , -1.4155476 , -0.78363633, -8.418715  ,
           4.5736566 ,  0.7427455 ,  3.7953236 , -0.08881438],
         [-1.9644134 , -3.1672404 , -0.80383766, -7.0728674 ,
           3.9525375 ,  2.5580957 ,  6.436445  , -8.407526  ],
         [-0.9375136 , -3.5294662 ,  0.27971148,  1.3742523 ,
          -1.3400316 , -0.73054576, -1.4859298 ,  2.189115  ],
         [-4.9197445 , -1.042948  ,  4.4484177 ,  6.059922  ,
           1.0065719 , -0.9886354 ,  1.507851  ,  4.1434894 ],
         [ 1.6555699 , -5.2372417 , -7.092151  ,  4.1866236 ,
          -3.8691092 ,  3.3705034 , -7.7330756 ,  4.9221206 ],
         [ 3.783893  , -1.7760451 , -6.8752346 , -1.0161141 ,
           5.096638  ,  6.8012614 , -0.33344686, -0.7447829 ]],

        [[ 7.382864  ,  1.568279  ,  8.594404  ,  8.14195   ,
          -5.8439884 , -6.345289  , -2.9872417 , -1.3761706 ],
         [ 8.14854   , -6.7695427 ,  4.0140567 ,  6.9059587 ,
           0.56577444, -7.546139  , -6.7947984 ,  1.3384073 ],
         [ 2.604012  ,  2.090725  ,  5.600357  ,  5.7823834 ,
          -7.842319  ,  4.2410126 ,  5.5899744 , -4.5454397 ],
         [-5.6148376 , -8.081939  , -8.879799  ,  4.702425  ,
           2.5586333 , -4.8888664 , -8.045462  , -8.366625  ],
         [ 8.165991  ,  2.9621363 ,  1.4714019 , -0.7440727 ,
           3.673987  ,  0.975381  ,  3.8852344 ,  2.2963078 ],
         [-4.636603  , -4.2065697 ,  2.3704987 ,  2.2805524 ,
           6.3544855 , -6.657332  ,  4.3716455 , -8.255994  ]]]],
      dtype=float32)}
Model:
<
   ir_version: 8,
   opset_import: ["" : 18, "onnxscript.atenlib" : 1, "this" : 1],
   producer_name: "pytorch",
   producer_version: "2.1.0"
>
torch_jit (float[4,3,8] input_0, float[4,4,6,8] input_1, float[4,4,6,8] input_2) => (float[4,4,3,8] 5) {
   3 = this._attention_scale (input_0)
   4 = this._causal_attention_mask (input_0, input_1)
   5 = onnxscript.atenlib._aten_scaled_dot_product_attention_float_mask_onnx <dropout_p = 0> (input_0, input_1, input_2, 4, 3)
}
<
  domain: "this",
  opset_import: ["" : 18]
>
_attention_scale (query) => (scale)
{
   tmp = Shape (query)
   subscript_index = Constant <value = int64[1] subscript_index {-1}> ()
   tmp_gather = Gather (tmp, subscript_index)
   subscript_axis = Constant <value = int64[1] subscript_axis {0}> ()
   tmp_0 = Squeeze (tmp_gather, subscript_axis)
   embedding_size = CastLike (tmp_0, query)
   tmp_1 = Constant <value = float tmp_1 {1}> ()
   tmp_2 = Sqrt (embedding_size)
   tmp_1_cast = CastLike (tmp_1, tmp_2)
   scale = Div (tmp_1_cast, tmp_2)
}
<
  domain: "this",
  opset_import: ["" : 18]
>
_causal_attention_mask (query, key) => (attn_mask_13)
{
   tmp = Shape (query)
   one = Constant <value = int64[1] one {1}> ()
   subscript_axis = Constant <value = int64[1] subscript_axis {0}> ()
   tmp_0 = Constant <value = int64 tmp_0 {-2}> ()
   tmp_0_reshaped = Reshape (tmp_0, one)
   tmp_1 = Constant <value = int64 tmp_1 {-1}> ()
   tmp_1_reshaped = Reshape (tmp_1, one)
   target_length = Slice (tmp, tmp_0_reshaped, tmp_1_reshaped, subscript_axis)
   tmp_2 = Shape (key)
   one_3 = Constant <value = int64[1] one_3 {1}> ()
   subscript_axis_4 = Constant <value = int64[1] subscript_axis_4 {0}> ()
   tmp_5 = Constant <value = int64 tmp_5 {-2}> ()
   tmp_5_reshaped = Reshape (tmp_5, one_3)
   tmp_6 = Constant <value = int64 tmp_6 {-1}> ()
   tmp_6_reshaped = Reshape (tmp_6, one_3)
   source_length = Slice (tmp_2, tmp_5_reshaped, tmp_6_reshaped, subscript_axis_4)
   size = Concat <axis = 0> (target_length, source_length)
   tmp_7 = Constant <value = float tmp_7 {1}> ()
   attn_mask = Expand (tmp_7, size)
   attn_mask_8 = Trilu <upper = 0> (attn_mask)
   tmp_9 = Constant <value = float tmp_9 {0}> ()
   tmp_9_cast = CastLike (tmp_9, attn_mask_8)
   tmp_10 = Equal (attn_mask_8, tmp_9_cast)
   tmp_11 = Constant <value_float = -inf> ()
   tmp_12 = Constant <value = float tmp_12 {0}> ()
   tmp_12_cast = CastLike (tmp_12, tmp_11)
   attn_mask_13 = Where (tmp_10, tmp_11, tmp_12_cast)
}
<
  domain: "onnxscript.atenlib",
  opset_import: ["" : 18]
>
_aten_scaled_dot_product_attention_float_mask_onnx <dropout_p>(query, key, value, attn_mask, scale) => (return_val)
{
   key_shape = Shape (key)
   one = Constant <value = int64[1] one {1}> ()
   subscript_axis = Constant <value = int64[1] subscript_axis {0}> ()
   tmp = Constant <value = int64 tmp {-1}> ()
   tmp_reshaped = Reshape (tmp, one)
   key_shape_shape = Shape (key_shape)
   key_shape_shape_dim = Gather (key_shape_shape, subscript_axis)
   key_last_dim = Slice (key_shape, tmp_reshaped, key_shape_shape_dim, subscript_axis)
   one_0 = Constant <value = int64[1] one_0 {1}> ()
   subscript_axis_1 = Constant <value = int64[1] subscript_axis_1 {0}> ()
   tmp_2 = Constant <value = int64 tmp_2 {-2}> ()
   tmp_2_reshaped = Reshape (tmp_2, one_0)
   tmp_3 = Constant <value = int64 tmp_3 {-1}> ()
   tmp_3_reshaped = Reshape (tmp_3, one_0)
   key_second_last_dim = Slice (key_shape, tmp_2_reshaped, tmp_3_reshaped, subscript_axis_1)
   one_4 = Constant <value = int64[1] one_4 {1}> ()
   subscript_axis_5 = Constant <value = int64[1] subscript_axis_5 {0}> ()
   tmp_6 = Constant <value = int64 tmp_6 {-2}> ()
   tmp_6_reshaped = Reshape (tmp_6, one_4)
   key_first_dims = Slice (key_shape, subscript_axis_5, tmp_6_reshaped, subscript_axis_5)
   tmp_7 = Constant <value_ints = [-1]> ()
   key_squeezed_shape = Concat <axis = 0> (tmp_7, key_second_last_dim, key_last_dim)
   key_squeezed = Reshape (key, key_squeezed_shape)
   key_squeezed_transposed = Transpose <perm = [0, 2, 1]> (key_squeezed)
   key_transposed_shape = Concat <axis = 0> (key_first_dims, key_last_dim, key_second_last_dim)
   key_transposed = Reshape (key_squeezed_transposed, key_transposed_shape)
   tmp_8 = Sqrt (scale)
   query_scaled = Mul (query, tmp_8)
   tmp_9 = Sqrt (scale)
   key_transposed_scaled = Mul (key_transposed, tmp_9)
   tmp_10 = MatMul (query_scaled, key_transposed_scaled)
   tmp_11 = Add (tmp_10, attn_mask)
   attn_weight = Softmax <axis = -1> (tmp_11)
   dropout_p = Constant <value_float: float = @dropout_p> ()
   attn_weight_12, _ = Dropout (attn_weight, dropout_p)
   return_val = MatMul (attn_weight_12, value)
}

[justinchuby]

https://github.com/microsoft/onnx-script/actions/runs/4635107849/jobs/8201882431

[justinchuby]

cc @BowenBao

[justinchuby]

Also https://github.com/microsoft/onnx-script/actions/runs/4639380134/jobs/8210161957#step:5:7619

[BowenBao]

Any idea why it is flaky?

Also it appears some other op tests are randomly failing?

[justinchuby]

Did you see others? So far I just saw attention failing

[justinchuby]

Any idea why it is flaky?

I am going to stress test it

[justinchuby]

WIP

import onnx
import onnxruntime as ort
import numpy as np



inputs = {'input_0': np.array([[[ 3.6824992 ,  0.50817776, -7.971771  , -3.3470182 ,
          5.651129  ,  5.1625967 ,  3.1133838 ,  7.9440765 ],
        [-4.679297  , -2.6819472 , -1.34092   , -2.211886  ,
          7.7628407 , -8.993065  ,  1.7052691 ,  8.833969  ],
        [-5.00294   ,  7.152945  , -7.754232  ,  1.5221719 ,
         -5.1234965 ,  2.8475125 ,  1.9688144 ,  1.9786946 ]],

       [[-2.0416298 , -2.787754  , -8.715644  ,  1.4930023 ,
          8.591377  , -6.7724156 , -3.462263  ,  2.3617163 ],
        [ 4.6034536 ,  3.2687244 ,  5.831397  ,  1.1378338 ,
         -2.2142708 , -7.502145  , -8.995697  , -6.89208   ],
        [ 4.521224  ,  8.088249  ,  8.566948  ,  6.740285  ,
          0.4644953 ,  6.197366  ,  3.4216266 , -6.0938964 ]],

       [[-1.6890653 , -8.352889  ,  7.725936  ,  1.9562176 ,
          8.264622  , -1.0898427 ,  7.146616  ,  4.2126226 ],
        [-1.0984998 , -0.29764366,  6.4143376 , -1.1836277 ,
          6.9980903 ,  1.974403  , -3.525754  , -4.171672  ],
        [-4.377688  ,  5.7745514 ,  6.803203  , -4.2311287 ,
         -8.21504   ,  5.101013  , -2.0465417 , -8.263691  ]],

       [[ 1.8895293 , -0.8955177 ,  6.484178  ,  1.9106255 ,
         -5.548896  , -1.6276009 , -6.462818  ,  4.208257  ],
        [ 2.6327426 ,  5.3127975 , -6.6609645 , -5.009738  ,
         -5.4769325 , -7.4629097 , -8.394707  , -7.2591786 ],
        [ 6.269464  , -2.715213  , -1.6603786 ,  2.7127583 ,
          2.3688498 ,  4.0463376 , -6.3547025 ,  7.9970865 ]]],
      dtype=np.float32),
 'input_1': np.array([[[[-7.1899724e-01, -4.2738266e+00,  5.8091240e+00,
          -7.4829779e+00, -5.3650808e+00,  5.0721526e+00,
           3.4888351e+00, -3.5642920e+00],
         [-7.1924722e-01, -1.6436620e+00,  3.6597099e+00,
           1.8056395e+00,  3.2981794e+00,  7.4339695e+00,
          -6.1898794e+00, -6.9231420e+00],
         [-2.1351576e+00, -4.0983467e+00, -6.4004545e+00,
           8.3525362e+00,  7.5186548e+00,  7.2669821e+00,
          -5.5431499e+00, -1.0642394e+00],
         [-7.0026522e+00, -3.2450545e+00,  7.4999037e+00,
           1.1283410e+00,  3.1221836e+00, -8.9466047e+00,
           3.7078857e+00,  6.0683627e+00],
         [-3.4560103e+00,  6.8391814e+00,  4.0536451e+00,
          -6.5714669e+00, -6.7465944e+00,  1.2635243e-01,
          -8.6807718e+00, -6.2895775e+00],
         [ 3.5202608e+00, -2.9611964e+00, -3.1890929e-01,
          -3.6772399e+00,  5.8250852e+00,  7.3904290e+00,
           6.4942417e+00, -8.0034370e+00]],

        [[-2.6608105e+00,  4.3589478e+00,  5.3230801e+00,
          -2.0171092e+00,  4.9280987e+00, -2.6609015e-01,
           9.9968612e-01,  6.0641217e+00],
         [ 6.1567974e+00, -6.3537207e+00, -5.4640074e+00,
           3.1882658e+00,  5.4091339e+00, -2.0955768e+00,
           7.9811831e+00, -1.0368057e+00],
         [-4.5640469e+00, -9.1866410e-01, -7.6602378e+00,
           2.7885296e+00, -8.2743692e+00,  8.3511810e+00,
          -6.8289633e+00,  5.8585949e+00],
         [-6.7507863e+00, -1.3116463e+00, -4.9842844e+00,
           4.5988898e+00,  5.2198944e+00, -8.4159822e+00,
          -7.7264547e+00,  1.1848015e+00],
         [-1.9361150e-01,  5.2183986e-02,  2.7360742e+00,
          -1.7787284e+00,  7.5948143e+00,  6.4017034e-01,
          -4.1461620e+00,  3.1228895e+00],
         [-6.9828029e+00, -4.3042555e+00,  2.6239922e+00,
          -3.7758765e+00, -8.5922298e+00,  4.5044374e+00,
           2.3163772e+00,  2.6195087e+00]],

        [[ 3.1282616e+00, -8.8809557e+00, -3.9208927e+00,
          -4.3719063e+00,  8.6117878e+00, -3.7917004e+00,
          -5.9125972e+00, -3.1587839e+00],
         [ 8.6476243e-01,  3.6269636e+00,  6.6279058e+00,
           6.4973092e+00,  9.2429030e-01,  8.1990528e+00,
           4.7852411e+00,  4.8019214e+00],
         [-7.7175350e+00,  7.3527718e-01, -1.9114966e+00,
          -8.5530787e+00, -4.4446745e+00,  6.8478780e+00,
           3.3004723e+00,  8.3176727e+00],
         [ 6.1366301e+00, -6.8849087e-02,  8.0888424e+00,
           1.8983184e+00, -1.3220093e+00, -3.8486180e+00,
           5.8009434e+00,  4.0348072e+00],
         [ 6.2488647e+00, -1.5634747e+00, -8.3295412e+00,
           1.6320630e+00,  9.8490608e-01, -5.8582330e+00,
           2.1819921e+00, -1.0892472e+00],
         [ 7.6784315e+00,  4.2354012e-01, -4.4449911e+00,
           7.2103410e+00,  2.6743126e-01, -1.5115514e+00,
           4.2986393e+00,  8.6815367e+00]],

        [[ 4.1037550e+00,  2.5377774e-01,  7.3711600e+00,
          -3.9469852e+00,  7.9539547e+00, -5.1278191e+00,
           1.6490865e-01,  8.1909714e+00],
         [ 3.5840235e+00, -1.8861337e+00,  5.5448837e+00,
           5.3087587e+00, -8.1125641e+00, -1.0030915e+00,
          -7.0829158e+00, -1.9399216e+00],
         [-1.0848806e+00, -3.6003990e+00, -3.2893891e+00,
           4.8675194e+00, -2.6728535e-01, -8.9547920e+00,
           1.5470392e+00, -8.8509493e+00],
         [-6.1742187e+00,  5.5013180e+00,  6.8800001e+00,
          -3.1932116e+00, -5.4649558e+00,  4.0727711e+00,
          -8.5974188e+00,  5.5524559e+00],
         [ 6.1993542e+00,  3.7977183e+00,  2.7468739e+00,
           1.8610615e+00, -6.8995171e+00, -3.6050830e+00,
           5.8053646e+00,  7.5802135e+00],
         [ 2.4039245e-01,  2.3122122e+00, -8.5365381e+00,
           6.3534451e+00, -5.2420683e+00, -4.8244820e+00,
           8.8806286e+00, -2.5799255e+00]]],


       [[[-4.5469117e+00,  1.2379907e+00,  7.0926232e+00,
           8.1777992e+00,  7.5339389e-01,  6.1271234e+00,
           6.4185357e+00, -1.5041968e+00],
         [-7.1115713e+00,  2.4213588e-01, -3.0572934e+00,
           3.1539569e+00, -6.5942035e+00, -6.0368500e+00,
          -6.7779636e+00, -6.7329454e+00],
         [ 6.8458652e+00, -8.9680138e+00, -3.3129036e-01,
          -5.5545340e+00, -3.7854156e+00, -2.5936584e+00,
           4.2222133e+00,  6.0517912e+00],
         [-2.0249572e+00,  2.9347765e-01, -9.8998189e-02,
           8.1959581e+00,  4.0224543e+00, -1.8686993e+00,
           6.0943651e+00, -7.2670989e+00],
         [ 3.5670311e+00,  6.9814262e+00, -4.6712914e+00,
           3.3676667e+00,  4.1266775e+00, -7.5507317e+00,
           2.5889549e+00,  6.8061590e+00],
         [ 8.2649374e-01, -1.7394587e+00, -3.5270941e+00,
          -5.9733334e+00,  2.1856475e+00, -8.5077410e+00,
          -7.7038922e+00, -3.2564583e+00]],

        [[-8.5227022e+00,  8.4039927e+00, -2.8999648e+00,
          -4.7849803e+00,  2.2623906e+00,  3.8732815e+00,
           6.7776151e+00, -3.6334953e+00],
         [-1.5726467e+00, -6.7962456e+00, -7.8587823e+00,
          -3.9159942e+00,  1.0126208e+00, -4.4616885e+00,
          -2.4844542e+00,  3.2988973e+00],
         [ 1.8569909e+00,  5.5531983e+00,  4.1560197e+00,
           4.7458696e+00, -4.6527410e+00,  4.8272972e+00,
          -8.0666418e+00,  2.7715511e+00],
         [-5.6761465e+00,  2.4851997e+00, -8.2341690e+00,
          -1.4906259e+00,  5.5357599e+00, -4.5007381e+00,
           5.5714383e+00, -1.3993331e+00],
         [ 8.4433994e+00,  8.4793863e+00, -5.3101525e+00,
          -2.1079812e+00, -1.8394611e+00, -3.4029670e+00,
           2.9584830e+00, -8.1945066e+00],
         [-5.3972964e+00, -1.4108773e+00, -8.0075903e+00,
          -8.6204233e+00,  3.6909633e+00, -7.5145659e+00,
           5.9796000e+00, -5.2258339e+00]],

        [[-7.1778381e-01,  5.4267430e+00, -8.8753090e+00,
          -8.7624407e+00,  3.7477648e-01, -3.6541085e+00,
           4.2068701e+00,  3.1413722e+00],
         [ 1.1077191e+00,  3.3783140e+00,  7.6997800e+00,
           1.4284147e+00,  5.9547205e+00, -5.0743222e+00,
           7.4154592e+00, -1.6003089e+00],
         [ 4.8724899e+00, -7.4795189e+00, -5.4084940e+00,
          -5.2685289e+00,  4.7543840e+00, -2.9990745e-01,
          -8.2597904e+00,  8.8265047e+00],
         [ 2.3457463e+00,  8.2777042e+00,  1.1234658e+00,
           1.2901769e+00,  6.4872026e+00, -6.1230860e+00,
           8.4994974e+00, -5.7886233e+00],
         [ 8.1574559e-01, -5.8533630e+00, -5.9631739e+00,
           6.9272938e+00, -2.3770370e+00, -1.6890750e+00,
           7.0318232e+00,  8.8172703e+00],
         [-8.1174850e+00, -1.5709591e-01,  2.5699918e+00,
           8.6411610e+00, -3.5880361e+00, -2.9290185e+00,
          -5.2750683e+00,  2.2996025e+00]],

        [[ 2.2196717e+00, -1.4192212e-01, -9.1280508e-01,
          -6.2656488e+00, -3.6966140e+00, -3.4228883e+00,
           2.3640594e+00,  3.9744167e+00],
         [ 4.2016773e+00, -6.8698926e+00, -6.9375014e-01,
          -1.7175858e+00,  8.9774923e+00,  6.4945574e+00,
          -6.7181258e+00, -4.9452114e+00],
         [-6.7987399e+00, -2.2524323e+00, -7.8746104e+00,
           4.2298632e+00, -3.9691062e+00, -4.5169935e+00,
           4.2535114e-01, -2.8756618e+00],
         [ 6.7990189e+00,  8.1583271e+00, -8.0357952e+00,
           4.0352116e+00,  1.8779433e-01,  4.7238102e+00,
          -3.3740537e+00, -6.8142700e-01],
         [-1.6835368e-01,  3.0858943e+00,  4.4873691e+00,
           3.5100384e+00,  1.1045508e+00, -7.1797361e+00,
          -5.1700535e+00, -5.1272292e+00],
         [-8.9690094e+00,  8.6211004e+00, -4.0746884e+00,
          -5.9822435e+00, -4.0757174e+00,  4.8867378e+00,
          -5.6418905e+00, -7.5201297e+00]]],


       [[[-6.7268429e+00,  1.9778920e+00, -3.6821258e-01,
          -8.1202641e+00, -6.1563187e+00, -7.7996106e+00,
           3.0874865e+00,  4.0729170e+00],
         [ 1.5330713e+00, -2.8834391e+00, -4.8220387e+00,
           3.5580716e+00, -2.0667257e+00, -1.4447011e+00,
          -2.9503708e+00,  1.5641710e+00],
         [ 2.1111710e+00,  7.3944802e+00,  7.9928842e+00,
          -2.3195195e+00,  1.3375951e+00, -6.0156112e+00,
          -6.7563653e+00,  7.4631319e+00],
         [ 1.0601882e+00, -4.4749684e+00,  4.5246806e+00,
           3.6308193e+00,  3.2216840e+00,  3.8081317e+00,
           4.8066580e-01, -6.5470457e+00],
         [-8.0008736e+00,  8.7865744e+00,  6.2830338e+00,
          -7.0215459e+00,  7.0528321e+00, -8.4174728e+00,
           5.9867773e+00,  6.2671132e+00],
         [ 5.6542921e+00,  3.8089805e+00, -5.4511852e+00,
          -3.8596795e+00,  1.9293430e+00,  3.5495539e+00,
           7.8894196e+00,  4.0034115e-01]],

        [[ 3.7061400e+00, -6.7969270e+00, -1.8267635e+00,
           5.3521914e+00,  7.3558080e-01,  7.1148372e+00,
          -7.2218165e+00, -6.0053234e+00],
         [ 2.8293808e+00, -3.7635362e-01, -7.6900253e+00,
           3.8401389e+00,  3.7094307e-01,  2.1764452e+00,
           1.9777976e+00, -6.2058625e+00],
         [-6.0370998e+00, -7.2009354e+00,  6.5371494e+00,
          -4.9404240e+00,  8.1020813e+00, -2.8741393e+00,
           1.9920863e+00,  8.5957775e+00],
         [-4.8297744e+00, -8.5363169e+00,  1.2978941e+00,
          -6.4928789e+00,  4.9591846e+00,  3.5780067e+00,
           6.5208597e+00, -9.3491828e-01],
         [-5.8799610e+00, -8.2938938e+00, -5.8130808e+00,
           8.0164347e+00, -1.9266254e+00, -1.6485178e-01,
          -1.9096481e+00, -4.5952625e+00],
         [ 7.3677340e+00, -2.3317699e+00, -4.3680472e+00,
           8.3155432e+00,  5.3867888e+00,  4.7965374e+00,
           6.4495573e+00, -5.7153873e+00]],

        [[ 2.9309044e+00, -2.5048571e+00,  3.2170415e-01,
           4.2393732e+00, -2.9874907e+00,  8.8308706e+00,
           4.4137359e-02,  5.1857948e+00],
         [ 8.2822859e-01, -8.6744375e+00, -5.3827901e+00,
          -7.4995098e+00,  1.0575017e+00, -6.0375986e+00,
          -7.1573448e+00,  1.9295082e+00],
         [-7.8780684e+00, -4.8974295e+00, -5.8221865e+00,
           4.3930702e+00,  8.2517281e+00,  3.1489263e+00,
          -5.0369902e+00,  6.8121619e+00],
         [ 8.3537655e+00, -2.4384992e+00,  6.6410818e+00,
           2.2752643e+00, -6.2602963e+00, -2.3506441e+00,
          -7.1778007e+00,  2.7155423e+00],
         [ 4.9769001e+00, -8.8521576e+00, -3.6783643e+00,
           1.9503844e-01,  7.2066655e+00, -4.9039664e+00,
          -5.2857006e-01, -2.5451589e-01],
         [-4.5259004e+00, -4.0339909e+00, -2.4080551e+00,
          -2.6892462e+00, -2.8968964e+00,  8.6046231e-01,
          -7.5238860e-01, -4.3800774e+00]],

        [[ 5.4502296e+00,  6.4385505e+00, -1.4108773e+00,
          -1.6989412e+00, -8.7716875e+00, -6.5796680e+00,
          -4.0990806e-01, -1.7158402e+00],
         [ 2.6331182e+00, -4.4771037e+00, -1.8640859e+00,
           1.3457490e+00,  8.8344049e+00,  7.0809565e+00,
          -5.9202008e+00, -7.4959481e-01],
         [-1.4734726e+00, -2.8769267e-01,  3.7493839e+00,
          -3.6205091e+00, -1.9184318e+00, -8.9154139e+00,
          -3.1971512e+00, -8.8269472e+00],
         [ 3.1860952e+00, -6.9278593e+00, -4.4472613e+00,
          -6.5265388e+00,  6.8438268e+00, -2.6420178e+00,
           7.8211799e+00, -4.4721718e+00],
         [ 6.8134761e+00, -1.7696207e+00, -6.7608910e+00,
           7.5262947e+00, -6.1228247e+00, -5.2924652e+00,
           4.9320297e+00, -8.4018159e-01],
         [-1.3669932e-01,  8.9240580e+00, -7.1422701e+00,
          -3.2507396e+00,  3.6613603e+00, -1.7330064e+00,
           1.5224218e-01, -3.0426936e+00]]],


       [[[-1.3942112e+00,  7.1839032e+00, -1.6729420e+00,
           3.6195810e+00,  6.7906194e+00,  6.9673986e+00,
           4.2673912e+00,  6.1863594e+00],
         [-2.3345013e+00,  7.1762502e-01, -3.8605795e+00,
          -8.7284765e+00,  7.5727806e+00, -8.6042204e+00,
          -5.6720405e+00, -3.2076559e+00],
         [-8.2325716e+00, -6.9842682e+00, -4.7076716e+00,
           1.0650001e+00, -7.8689532e+00, -6.4693575e+00,
          -5.2609997e+00, -6.7677979e+00],
         [ 2.1124778e+00,  3.0450153e-01,  1.9749856e+00,
           8.1192884e+00, -3.9406061e+00, -3.3742983e+00,
           2.8711147e+00,  7.8933926e+00],
         [ 2.8725107e+00, -7.7150846e-01,  2.9235039e+00,
          -8.8649769e+00, -6.1562929e+00, -3.7844102e+00,
          -8.1471882e+00, -6.7970886e+00],
         [-6.9216795e+00,  2.5836787e+00,  2.7266221e+00,
          -2.1676862e-01, -2.0801926e+00,  2.4266374e-01,
           5.4382429e+00, -6.7768016e+00]],

        [[ 6.3940697e+00,  3.9143558e+00, -3.1416392e+00,
           6.3145533e+00, -7.4543953e-03, -4.1244521e+00,
           8.1106386e+00,  2.6600432e+00],
         [ 6.7920012e+00, -2.4422779e+00, -7.7121687e+00,
           7.2174406e+00,  5.2065647e-01,  4.2741752e+00,
          -2.3327858e+00,  4.2567835e+00],
         [-7.1112967e+00,  3.1746733e-01,  6.8336978e+00,
          -2.4167197e+00,  7.7300525e+00,  3.3097420e+00,
           4.4366388e+00,  8.8930531e+00],
         [-2.1897941e+00, -2.0952549e+00,  5.6594954e+00,
          -4.5546107e+00, -6.6358676e+00,  5.8684425e+00,
          -7.8526454e+00, -4.2624216e+00],
         [ 2.1807680e+00,  1.3253288e+00, -5.4561248e+00,
          -2.2769079e+00,  5.7771726e+00, -5.8317447e+00,
           5.1623950e+00, -9.6378958e-01],
         [-3.2215466e+00, -8.5734272e+00, -2.3606281e+00,
           7.4004955e+00, -4.9304538e+00, -7.1427946e+00,
          -7.0218408e-01,  6.3800020e+00]],

        [[-2.4413552e+00,  8.8226805e+00,  8.6688747e+00,
           8.2456303e+00, -5.1477995e+00, -6.4066162e+00,
          -8.4346704e+00, -1.7500427e+00],
         [-5.8611622e+00,  4.1928344e+00,  8.1168687e-01,
           7.1380105e+00, -8.3490257e+00, -7.5629129e+00,
          -1.6748067e+00,  2.3720598e+00],
         [ 7.5774841e+00, -7.5245619e-01,  7.6912527e+00,
          -2.4364393e+00, -1.2704905e+00, -3.4654548e+00,
          -3.8250360e+00, -3.4094546e+00],
         [-7.7797914e+00, -1.1047976e+00, -3.0346942e+00,
          -5.8269124e+00, -6.6757545e+00, -8.6790047e+00,
          -2.0167615e+00, -7.3273220e+00],
         [ 8.0918541e+00, -2.4159055e+00,  2.4191928e+00,
          -6.4524441e+00,  9.1833794e-01, -8.4585638e+00,
           4.0517311e+00,  4.2503734e+00],
         [-8.7459154e+00,  3.9245460e+00,  1.1458172e+00,
          -5.6210508e+00, -4.1205220e+00,  8.8462982e+00,
          -4.6813412e+00,  7.8744593e+00]],

        [[ 1.8111606e+00,  4.1234994e+00, -4.5300570e+00,
           2.9255669e+00,  1.9378874e+00,  7.1779141e+00,
          -8.8999577e+00, -3.2721491e+00],
         [ 2.2411036e-01, -6.4597197e+00,  1.5338030e+00,
           5.8029237e+00,  5.4236789e+00, -4.6408010e+00,
           3.6380281e+00, -6.5486774e+00],
         [ 5.5946975e+00, -7.3559809e+00, -8.1593246e+00,
           3.1234422e+00, -8.4961116e-01,  2.2184992e+00,
           8.3744087e+00,  8.6216440e+00],
         [ 5.0746579e+00, -8.1641111e+00, -1.3856753e+00,
           9.8382568e-01, -8.7636156e+00,  7.5265608e+00,
           5.9059334e+00, -4.1574841e+00],
         [-1.8782684e+00, -5.2204351e+00,  3.4616139e+00,
           1.5521868e+00, -3.1202180e+00,  4.8116479e+00,
          -1.9454974e+00,  9.6651042e-01],
         [ 6.2313490e+00, -4.2285423e+00, -8.3006687e+00,
           4.8575330e+00, -5.8817658e+00,  3.4023521e+00,
          -2.7057729e+00,  6.0415225e+00]]]], dtype=np.float32),
 'input_2': np.array([[[[-0.34230995, -4.0978384 , -6.349059  ,  6.047771  ,
          -5.631379  ,  2.551255  ,  6.882798  , -7.2295017 ],
         [-7.553694  ,  0.91815555, -3.5329392 , -6.606769  ,
           5.9512577 , -8.372425  , -4.1034985 , -7.4758034 ],
         [ 2.3787193 ,  0.7762345 ,  3.754082  ,  3.768081  ,
          -0.0356015 , -2.702716  ,  7.9581776 ,  0.3742733 ],
         [ 2.1427407 , -7.4108696 , -8.820838  , -4.395982  ,
          -4.7160378 ,  2.6075492 , -4.000244  ,  0.93990827],
         [ 8.340205  , -2.6548986 ,  2.2993689 ,  0.11037505,
           5.922707  ,  3.936295  , -0.41812956,  4.922632  ],
         [-0.45790136, -1.5300007 , -3.5081182 , -7.342042  ,
           7.5179358 , -5.703294  , -8.579783  , -4.036794  ]],

        [[ 1.8180141 , -0.81675303,  4.0971346 ,  5.156386  ,
           5.2929974 , -7.5512433 ,  3.6692598 ,  4.4550023 ],
         [-4.9928417 ,  0.04359448,  1.3489387 ,  1.1285588 ,
           4.8295083 ,  7.6410804 ,  8.644264  ,  2.0150535 ],
         [ 5.263167  ,  3.8479776 , -4.8689647 ,  6.0192933 ,
          -0.63945794,  3.672884  , -5.973283  ,  7.200478  ],
         [ 2.008657  ,  2.0394938 , -8.0063925 , -0.78217936,
          -3.147153  ,  8.497046  , -8.23968   , -2.220986  ],
         [-4.4639134 , -6.524279  ,  3.4615722 , -6.213848  ,
           4.254297  ,  1.0802768 , -4.764764  , -1.024209  ],
         [ 0.7804874 , -4.010984  , -5.782582  , -5.071845  ,
           2.171483  , -7.930274  ,  5.6670957 ,  5.9107027 ]],

        [[ 8.275736  ,  7.6145887 , -5.0580616 , -3.905322  ,
           6.5873466 , -1.164304  , -6.626468  , -4.31118   ],
         [ 6.31756   , -5.1179714 ,  3.165544  , -7.5146894 ,
           5.9429803 ,  2.475356  ,  6.2080154 ,  0.306952  ],
         [ 4.568535  ,  7.375328  , -1.1353394 , -6.80179   ,
          -6.3504    ,  5.276816  , -5.7391677 , -7.878265  ],
         [-7.620439  , -3.9857883 ,  1.6696558 , -8.154997  ,
          -4.170337  ,  8.616638  , -1.5307474 , -3.037105  ],
         [ 4.2566166 , -7.4061036 , -7.0274777 ,  1.2868155 ,
           2.8450866 , -8.869081  ,  7.9074764 , -3.504482  ],
         [ 1.3749079 ,  6.397666  ,  3.4449348 ,  2.7945228 ,
          -8.211373  , -8.39686   ,  3.3470879 ,  0.3344618 ]],

        [[ 5.526359  ,  4.50389   , -7.2203007 ,  1.6565012 ,
          -7.0605087 , -4.975044  ,  2.4733067 , -0.503073  ],
         [-5.747767  ,  3.0517874 ,  6.24559   , -8.845017  ,
           8.441818  ,  4.968251  , -0.7585845 , -8.727442  ],
         [-3.8099234 ,  2.746667  , -3.834444  , -2.958364  ,
          -3.7167199 ,  7.738069  ,  7.6167045 , -7.0197434 ],
         [-2.7766004 ,  8.151881  , -0.5183884 , -4.3818927 ,
           1.0883235 ,  7.1210847 , -1.1453065 , -8.492519  ],
         [-3.0083947 ,  4.479186  , -1.3297963 ,  6.358906  ,
           3.0952597 ,  0.0789299 ,  3.437828  , -1.2148347 ],
         [-7.7152843 , -7.713045  , -7.0792937 , -1.4334733 ,
          -6.5075026 ,  1.3189505 ,  1.0319617 ,  6.8479414 ]]],


       [[[ 0.5060556 ,  1.5339446 ,  3.2002883 , -2.544455  ,
          -4.484561  , -8.466113  , -8.765525  ,  0.59179723],
         [-4.2559876 ,  2.4915285 ,  2.2367692 ,  7.0193615 ,
          -1.0987122 , -0.0380584 ,  6.211625  ,  1.5262123 ],
         [ 4.119171  , -4.8599524 , -6.5050607 ,  6.861494  ,
           6.415559  , -3.4077444 ,  1.4950504 , -8.004093  ],
         [ 0.629861  , -8.355907  ,  5.086895  , -8.920359  ,
          -2.1710293 ,  8.396976  , -3.0490515 , -5.216814  ],
         [ 8.11894   , -4.5523977 ,  1.9849752 ,  5.029893  ,
           2.418108  , -6.2179594 ,  8.168247  ,  4.3157053 ],
         [-5.65954   , -5.1365924 , -2.0757604 ,  3.0025828 ,
           3.4030323 ,  1.1588559 ,  5.9023085 ,  8.953258  ]],

        [[ 1.8515954 ,  3.525226  ,  5.118249  ,  1.7841175 ,
           0.5548954 ,  8.092021  ,  6.2133555 ,  5.8781395 ],
         [-3.5019822 , -4.3029294 , -4.880156  , -7.5367327 ,
          -0.91440153,  5.070908  ,  1.3185772 , -3.1836286 ],
         [-1.2450696 , -6.896912  , -5.781066  ,  0.82857835,
          -5.229164  ,  1.1092318 ,  7.422377  ,  2.30473   ],
         [ 6.889682  , -5.7019978 ,  0.23406887, -8.8410225 ,
           0.06235921, -6.630025  ,  8.446856  , -7.6939692 ],
         [ 3.07654   , -5.3836336 ,  2.491255  , -0.04694295,
          -1.151144  , -3.4929657 ,  5.7868576 , -1.62457   ],
         [ 0.767115  , -2.7122412 , -6.558548  , -3.7508407 ,
           8.069345  ,  8.56679   , -4.181401  ,  0.89693284]],

        [[ 1.4608608 ,  7.9395723 ,  1.2559626 ,  0.41896963,
          -2.6179614 ,  1.8980459 , -8.961205  , -8.979448  ],
         [-1.0432377 , -6.242316  , -6.250534  , -1.5276393 ,
           4.9933357 ,  4.4542627 , -3.1269214 , -0.6360065 ],
         [ 8.355031  ,  1.9468868 ,  6.120341  , -3.7278447 ,
          -8.053922  , -4.172763  ,  0.7909266 , -2.1797915 ],
         [-4.914157  ,  8.654715  , -8.307294  ,  0.22850919,
           3.586665  , -2.4280643 ,  3.9034038 ,  0.03194726],
         [-5.3498836 , -4.8937836 , -4.04868   , -8.787451  ,
           5.609341  ,  8.249806  , -4.623246  ,  0.4448465 ],
         [ 3.007388  ,  2.5482736 , -4.995444  , -5.4690695 ,
           8.761045  ,  5.240899  ,  8.413323  ,  6.3179026 ]],

        [[-0.6696371 ,  6.336029  , -0.28866792, -0.0558393 ,
          -3.8995607 ,  7.0187025 , -1.9134021 ,  1.3067003 ],
         [-7.176584  ,  3.0873265 , -4.2791715 , -2.0322926 ,
          -0.8047968 ,  3.3364782 ,  0.3673811 ,  1.56819   ],
         [-5.8456073 , -6.80033   ,  3.2412179 , -3.706018  ,
          -7.759484  , -3.714308  , -8.829413  ,  1.7982377 ],
         [ 0.6684011 ,  8.311016  , -4.962897  , -4.6545744 ,
           2.1489978 , -2.3648608 ,  7.0137653 ,  1.3691282 ],
         [-0.19555664, -3.4528592 ,  5.4378805 , -4.212881  ,
           3.6190457 ,  5.7816677 , -7.059673  , -7.9743114 ],
         [-5.075494  ,  4.318495  , -4.964339  ,  0.8507484 ,
           1.3588833 ,  1.8894832 , -4.589383  ,  0.3146907 ]]],


       [[[ 6.2364674 ,  4.665291  ,  1.7193228 , -7.157478  ,
           4.2342606 , -3.583723  ,  6.2986712 ,  3.7378986 ],
         [-1.6520036 ,  3.4791384 , -3.7425292 ,  1.5534518 ,
          -0.44907153,  0.14281583,  3.1689956 ,  8.764865  ],
         [-3.9725296 ,  8.788686  , -1.6515831 ,  8.559981  ,
           2.0253005 ,  0.7339339 , -3.0099182 , -4.742607  ],
         [-8.850726  , -4.9126387 , -0.80580103,  7.571518  ,
           4.8703218 ,  1.5691963 , -8.71032   ,  6.340592  ],
         [ 6.196273  , -4.0458193 ,  3.6451972 , -7.4987164 ,
          -4.5503087 , -2.716004  , -8.160566  , -0.5498121 ],
         [ 1.8430274 ,  2.6196856 ,  8.422474  ,  5.1977305 ,
          -8.234614  , -2.8249798 ,  5.042182  , -8.638262  ]],

        [[-0.5566045 ,  3.4055567 , -6.0347867 ,  7.7202396 ,
          -5.933469  , -6.6021404 , -6.831831  , -1.5764834 ],
         [ 3.0054526 ,  2.5924494 , -4.446918  ,  7.871669  ,
          -4.823057  ,  7.061883  ,  3.023122  , -2.904017  ],
         [ 6.656292  ,  2.9841828 , -7.652502  ,  4.5741286 ,
           8.511661  , -6.7322598 , -1.897414  , -7.2625446 ],
         [ 6.033968  ,  4.775005  ,  6.5643225 ,  2.5277867 ,
          -1.4840738 , -0.8781692 , -6.7057233 ,  5.7115536 ],
         [-3.6106932 ,  0.9877299 ,  0.30608404,  2.7207212 ,
          -7.229786  , -1.4255319 ,  4.15208   , -6.652448  ],
         [-3.866281  , -4.9557824 , -7.6150084 ,  8.787146  ,
           7.7758656 ,  5.639732  , -7.460018  , -0.08213997]],

        [[ 1.6527236 ,  3.9922428 ,  1.8127184 , -5.516509  ,
           0.2853613 , -7.2778115 ,  4.5133758 ,  5.669015  ],
         [-3.2203126 , -7.8779235 ,  0.08367848,  4.6010675 ,
          -6.14085   ,  3.799644  ,  6.268093  , -7.460702  ],
         [ 4.977277  ,  4.0702305 , -4.6315937 ,  7.553576  ,
          -4.3281403 ,  6.840512  , -6.898634  ,  7.098889  ],
         [-2.0235925 , -7.9501467 ,  4.7939253 , -0.73716974,
           2.5420046 ,  4.1460257 ,  0.9609475 ,  2.8745513 ],
         [-7.585773  , -7.9436517 , -4.2263126 , -4.7055035 ,
          -8.483265  ,  1.4052114 , -0.7140695 , -1.9820827 ],
         [-0.4182197 ,  4.3249044 , -3.3396592 ,  0.56857896,
          -0.8039224 ,  4.3460445 , -0.9574542 ,  3.2329597 ]],

        [[-3.523854  ,  6.971783  ,  1.5148151 ,  1.5140405 ,
           4.357374  , -8.971563  ,  3.9768887 ,  7.230226  ],
         [ 5.390065  ,  5.844248  , -7.2293396 ,  7.7898884 ,
           2.0122254 , -8.772329  ,  8.143966  ,  7.018045  ],
         [-1.4358433 ,  8.825392  ,  0.5666628 ,  7.2184963 ,
          -7.9558535 , -8.390788  ,  2.5075521 , -0.06732452],
         [-7.6399775 , -5.0466223 ,  5.0448103 ,  5.7374845 ,
          -6.8261423 , -4.0237017 ,  0.05233955, -0.01148844],
         [-6.628728  ,  6.655866  , -6.2843213 ,  6.4027185 ,
          -5.566676  , -8.942067  ,  2.2423835 , -4.305365  ],
         [ 6.067338  , -2.9838877 ,  0.8593594 ,  6.6992774 ,
          -5.8451815 , -3.7053301 ,  6.127894  ,  7.5919476 ]]],


       [[[ 1.5686717 ,  6.9187164 ,  5.444451  ,  5.0889187 ,
          -0.35480368,  8.689396  , -2.861092  ,  0.612609  ],
         [-2.7166636 , -0.6744189 ,  5.4085417 ,  6.823352  ,
           1.3416077 ,  5.916538  ,  6.3087487 ,  4.37452   ],
         [-2.4527354 ,  5.251175  ,  7.5547123 ,  7.4452276 ,
           2.331995  ,  6.6418686 , -6.217531  ,  1.7046597 ],
         [-7.907222  ,  4.620072  ,  0.41317284, -1.731353  ,
           5.737301  , -1.3776695 ,  3.6747088 ,  4.8420577 ],
         [-4.46467   , -7.796344  ,  0.8038162 , -5.5036    ,
          -5.1038737 ,  4.238401  ,  2.138669  , -2.0001135 ],
         [-3.271941  , -6.6557193 , -3.8459122 ,  0.57687986,
          -5.2578635 , -8.170835  , -8.719526  , -8.48733   ]],

        [[-4.1398983 ,  2.4248843 , -4.115657  ,  5.945626  ,
           4.9255066 , -7.0417085 ,  2.215833  , -3.8760872 ],
         [ 7.4642477 ,  6.81977   ,  4.1280193 , -4.5070915 ,
           1.3013864 , -5.674239  ,  1.3475128 ,  8.897213  ],
         [-1.7353356 , -5.1756988 , -4.387149  , -5.5859427 ,
           4.559204  , -5.3083544 ,  6.2997246 ,  5.254262  ],
         [ 5.4144382 , -0.43160605, -1.8787909 ,  3.2602606 ,
          -1.3443532 ,  3.6168933 ,  2.0089831 ,  3.1299973 ],
         [-3.8472211 ,  5.579997  , -1.7262236 ,  5.113995  ,
          -0.5548074 , -8.845102  , -0.7667781 , -2.8695989 ],
         [-4.4622483 , -0.3521011 , -3.4386177 ,  0.70515597,
          -2.6571198 , -5.584134  ,  3.3040287 , -7.370762  ]],

        [[-6.967508  , -1.4155476 , -0.78363633, -8.418715  ,
           4.5736566 ,  0.7427455 ,  3.7953236 , -0.08881438],
         [-1.9644134 , -3.1672404 , -0.80383766, -7.0728674 ,
           3.9525375 ,  2.5580957 ,  6.436445  , -8.407526  ],
         [-0.9375136 , -3.5294662 ,  0.27971148,  1.3742523 ,
          -1.3400316 , -0.73054576, -1.4859298 ,  2.189115  ],
         [-4.9197445 , -1.042948  ,  4.4484177 ,  6.059922  ,
           1.0065719 , -0.9886354 ,  1.507851  ,  4.1434894 ],
         [ 1.6555699 , -5.2372417 , -7.092151  ,  4.1866236 ,
          -3.8691092 ,  3.3705034 , -7.7330756 ,  4.9221206 ],
         [ 3.783893  , -1.7760451 , -6.8752346 , -1.0161141 ,
           5.096638  ,  6.8012614 , -0.33344686, -0.7447829 ]],

        [[ 7.382864  ,  1.568279  ,  8.594404  ,  8.14195   ,
          -5.8439884 , -6.345289  , -2.9872417 , -1.3761706 ],
         [ 8.14854   , -6.7695427 ,  4.0140567 ,  6.9059587 ,
           0.56577444, -7.546139  , -6.7947984 ,  1.3384073 ],
         [ 2.604012  ,  2.090725  ,  5.600357  ,  5.7823834 ,
          -7.842319  ,  4.2410126 ,  5.5899744 , -4.5454397 ],
         [-5.6148376 , -8.081939  , -8.879799  ,  4.702425  ,
           2.5586333 , -4.8888664 , -8.045462  , -8.366625  ],
         [ 8.165991  ,  2.9621363 ,  1.4714019 , -0.7440727 ,
           3.673987  ,  0.975381  ,  3.8852344 ,  2.2963078 ],
         [-4.636603  , -4.2065697 ,  2.3704987 ,  2.2805524 ,
           6.3544855 , -6.657332  ,  4.3716455 , -8.255994  ]]]],
      dtype=np.float32)}

model_text = """
<
ir_version: 8,
opset_import: ["" : 18, "onnxscript.atenlib" : 1, "this" : 1],
producer_name: "pytorch",
producer_version: "2.1.0"
>
torch_jit (float[4,3,8] input_0, float[4,4,6,8] input_1, float[4,4,6,8] input_2) => (float[4,4,3,8] output) {
intermediate_3 = this._attention_scale (input_0)
intermediate_4 = this._causal_attention_mask (input_0, input_1)
output = onnxscript.atenlib._aten_scaled_dot_product_attention_float_mask_onnx <dropout_p = 0> (input_0, input_1, input_2, intermediate_4, intermediate_3)
}
<
domain: "this",
opset_import: ["" : 18]
>
_attention_scale (query) => (scale)
{
tmp = Shape (query)
subscript_index = Constant <value = int64[1] subscript_index {-1}> ()
tmp_gather = Gather (tmp, subscript_index)
subscript_axis = Constant <value = int64[1] subscript_axis {0}> ()
tmp_0 = Squeeze (tmp_gather, subscript_axis)
embedding_size = CastLike (tmp_0, query)
tmp_1 = Constant <value = float tmp_1 {1}> ()
tmp_2 = Sqrt (embedding_size)
tmp_1_cast = CastLike (tmp_1, tmp_2)
scale = Div (tmp_1_cast, tmp_2)
}
<
domain: "this",
opset_import: ["" : 18]
>
_causal_attention_mask (query, key) => (attn_mask_13)
{
tmp = Shape (query)
one = Constant <value = int64[1] one {1}> ()
subscript_axis = Constant <value = int64[1] subscript_axis {0}> ()
tmp_0 = Constant <value = int64 tmp_0 {-2}> ()
tmp_0_reshaped = Reshape (tmp_0, one)
tmp_1 = Constant <value = int64 tmp_1 {-1}> ()
tmp_1_reshaped = Reshape (tmp_1, one)
target_length = Slice (tmp, tmp_0_reshaped, tmp_1_reshaped, subscript_axis)
tmp_2 = Shape (key)
one_3 = Constant <value = int64[1] one_3 {1}> ()
subscript_axis_4 = Constant <value = int64[1] subscript_axis_4 {0}> ()
tmp_5 = Constant <value = int64 tmp_5 {-2}> ()
tmp_5_reshaped = Reshape (tmp_5, one_3)
tmp_6 = Constant <value = int64 tmp_6 {-1}> ()
tmp_6_reshaped = Reshape (tmp_6, one_3)
source_length = Slice (tmp_2, tmp_5_reshaped, tmp_6_reshaped, subscript_axis_4)
size = Concat <axis = 0> (target_length, source_length)
tmp_7 = Constant <value = float tmp_7 {1}> ()
attn_mask = Expand (tmp_7, size)
attn_mask_8 = Trilu <upper = 0> (attn_mask)
tmp_9 = Constant <value = float tmp_9 {0}> ()
tmp_9_cast = CastLike (tmp_9, attn_mask_8)
tmp_10 = Equal (attn_mask_8, tmp_9_cast)
tmp_11 = Constant <value_float = -inf> ()
tmp_12 = Constant <value = float tmp_12 {0}> ()
tmp_12_cast = CastLike (tmp_12, tmp_11)
attn_mask_13 = Where (tmp_10, tmp_11, tmp_12_cast)
}
<
domain: "onnxscript.atenlib",
opset_import: ["" : 18]
>
_aten_scaled_dot_product_attention_float_mask_onnx <dropout_p>(query, key, value, attn_mask, scale) => (return_val)
{
key_shape = Shape (key)
one = Constant <value = int64[1] one {1}> ()
subscript_axis = Constant <value = int64[1] subscript_axis {0}> ()
tmp = Constant <value = int64 tmp {-1}> ()
tmp_reshaped = Reshape (tmp, one)
key_shape_shape = Shape (key_shape)
key_shape_shape_dim = Gather (key_shape_shape, subscript_axis)
key_last_dim = Slice (key_shape, tmp_reshaped, key_shape_shape_dim, subscript_axis)
one_0 = Constant <value = int64[1] one_0 {1}> ()
subscript_axis_1 = Constant <value = int64[1] subscript_axis_1 {0}> ()
tmp_2 = Constant <value = int64 tmp_2 {-2}> ()
tmp_2_reshaped = Reshape (tmp_2, one_0)
tmp_3 = Constant <value = int64 tmp_3 {-1}> ()
tmp_3_reshaped = Reshape (tmp_3, one_0)
key_second_last_dim = Slice (key_shape, tmp_2_reshaped, tmp_3_reshaped, subscript_axis_1)
one_4 = Constant <value = int64[1] one_4 {1}> ()
subscript_axis_5 = Constant <value = int64[1] subscript_axis_5 {0}> ()
tmp_6 = Constant <value = int64 tmp_6 {-2}> ()
tmp_6_reshaped = Reshape (tmp_6, one_4)
key_first_dims = Slice (key_shape, subscript_axis_5, tmp_6_reshaped, subscript_axis_5)
tmp_7 = Constant <value_ints = [-1]> ()
key_squeezed_shape = Concat <axis = 0> (tmp_7, key_second_last_dim, key_last_dim)
key_squeezed = Reshape (key, key_squeezed_shape)
key_squeezed_transposed = Transpose <perm = [0, 2, 1]> (key_squeezed)
key_transposed_shape = Concat <axis = 0> (key_first_dims, key_last_dim, key_second_last_dim)
key_transposed = Reshape (key_squeezed_transposed, key_transposed_shape)
tmp_8 = Sqrt (scale)
query_scaled = Mul (query, tmp_8)
tmp_9 = Sqrt (scale)
key_transposed_scaled = Mul (key_transposed, tmp_9)
tmp_10 = MatMul (query_scaled, key_transposed_scaled)
tmp_11 = Add (tmp_10, attn_mask)
attn_weight = Softmax <axis = -1> (tmp_11)
dropout_p = Constant <value_float: float = @dropout_p> ()
attn_weight_12, _ = Dropout (attn_weight, dropout_p)
return_val = MatMul (attn_weight_12, value)
}
"""

model = onnx.parser.parse_model(model_text)

model = onnx.shape_inference.infer_shapes(model)
onnx.checker.check_model(model, full_check=True)

serialized = model.SerializeToString()
session = ort.InferenceSession(serialized, providers=['CPUExecutionProvider'])
# Run once
session.run(None, inputs)
print("done")

[BowenBao]

nvm the other one is matmul in the first link. I suppose that is resolved.