Skip to content

Add pass to convert atomic rmw to non-atomic ops when legal #1782

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
ivanradanov wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add pass to convert atomic rmw to non-atomic ops when legal #1782

ivanradanov wants to merge 2 commits into from

Conversation

@ivanradanov
Copy link
Collaborator

@ivanradanov ivanradanov commented Dec 16, 2025

No description provided.

@ivanradanov ivanradanov requested a review from wsmoses December 16, 2025 14:26
github-actions[bot]
github-actions bot reviewed Dec 16, 2025
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

EnzymeJAX Benchmarks

Details
Benchmark suite Current: c1992d4 Previous: 4a1165c Ratio
actmtch / JaXPipe / cpu / Primal 0.000008764808000705671 s 0.000005010517999835429 s 1.75
actmtch / Jax / cpu / Primal 0.000008885766001185401 s 0.000004809449999811477 s 1.85
actmtch / HLOOpt / cpu / Primal 0.00001089175600282033 s 0.00000530116799973257 s 2.05
actmtch / PartOpt / cpu / Primal 0.00000960414499786566 s 0.000004738015999919299 s 2.03
actmtch / IPartOpt / cpu / Primal 0.000009239004997652956 s 0.000004711642000074789 s 1.96
actmtch / DefOpt / cpu / Primal 0.000010957680002320556 s 0.0000053095960001883215 s 2.06
actmtch / IDefOpt / cpu / Primal 0.000010912956000538544 s 0.00000531172199998764 s 2.05
actmtch / JaXPipe / cpu / Forward 0.000016766919998190133 s 0.000008042900999953417 s 2.08
actmtch / Jax / cpu / Forward 0.000012844130000303268 s 0.0000071996839997154895 s 1.78
actmtch / HLOOpt / cpu / Forward 0.00001586150399816688 s 0.000008179603999906248 s 1.94
actmtch / PartOpt / cpu / Forward 0.00001692909499979578 s 0.000008225519999996323 s 2.06
actmtch / IPartOpt / cpu / Forward 0.000016936281001108 s 0.000008223061000080633 s 2.06
actmtch / DefOpt / cpu / Forward 0.000016611181999905968 s 0.000008229308999943897 s 2.02
actmtch / IDefOpt / cpu / Forward 0.000016754129999753786 s 0.000008161577000009856 s 2.05
actmtch / JaXPipe / cpu / PreRev 0.000017084702001739062 s 0.000008187203000034061 s 2.09
actmtch / JaXPipe / cpu / PostRev 0.000013019424997764873 s 0.000007099333000041952 s 1.83
actmtch / JaXPipe / cpu / BothRev 0.000016774617000919532 s 0.000008179985999959171 s 2.05
actmtch / Jax / cpu / BothRev 0.000013064210997981718 s 0.000007223351999982696 s 1.81
actmtch / HLOOpt / cpu / PreRev 0.00001670318699689233 s 0.00000831829100025061 s 2.01
actmtch / HLOOpt / cpu / PostRev 0.000015985330999683354 s 0.000008230335000007471 s 1.94
actmtch / HLOOpt / cpu / BothRev 0.000016698704999726033 s 0.000008109762999993109 s 2.06
actmtch / PartOpt / cpu / PreRev 0.000015962781999405708 s 0.000008014170000024023 s 1.99
actmtch / PartOpt / cpu / PostRev 0.000013697408998268656 s 0.000007239625000238447 s 1.89
actmtch / PartOpt / cpu / BothRev 0.000016080976998637196 s 0.000008151960999839503 s 1.97
actmtch / IPartOpt / cpu / PreRev 0.000016893321997486054 s 0.000008258112000021356 s 2.05
actmtch / IPartOpt / cpu / PostRev 0.00001349494600071921 s 0.000007345655000335683 s 1.84
actmtch / IPartOpt / cpu / BothRev 0.000016670946999511216 s 0.000008073200000126234 s 2.06
actmtch / DefOpt / cpu / PreRev 0.000016048045999923488 s 0.000008086352000191255 s 1.98
actmtch / DefOpt / cpu / PostRev 0.00001687758399930317 s 0.000008116398999845841 s 2.08
actmtch / DefOpt / cpu / BothRev 0.000016867709997313796 s 0.000008047573000112606 s 2.10
actmtch / IDefOpt / cpu / PreRev 0.000016862656997545854 s 0.000008090865000212944 s 2.08
actmtch / IDefOpt / cpu / PostRev 0.000016756209002778633 s 0.000008463122000193834 s 1.98
actmtch / IDefOpt / cpu / BothRev 0.00001691553899945575 s 0.000008146463999764818 s 2.08
actmtch / JaXPipe / tpu / Primal 0.0001539747530005 s 0.0001535154900011 s 1.00
actmtch / Jax / tpu / Primal 0.0001529509840001 s 0.0001530672299995 s 1.00
actmtch / HLOOpt / tpu / Primal 0.0001535086129988 s 0.0001544030010009 s 0.99
actmtch / PartOpt / tpu / Primal 0.000151818803999 s 0.0001377300819985 s 1.10
actmtch / IPartOpt / tpu / Primal 0.0001512362439971 s 0.0001389470540016 s 1.09
actmtch / DefOpt / tpu / Primal 0.0001394300880019 s 0.0001400444040009 s 1.00
actmtch / IDefOpt / tpu / Primal 0.0001314456910004 s 0.0001380777029989 s 0.95
actmtch / JaXPipe / tpu / Forward 0.0001990186770017 s 0.0002225152820028 s 0.89
actmtch / Jax / tpu / Forward 0.000220320168999 s 0.0002306253950009 s 0.96
actmtch / HLOOpt / tpu / Forward 0.0002315494739996 s 0.000223667981998 s 1.04
actmtch / PartOpt / tpu / Forward 0.0002228471970011 s 0.0002129827069984 s 1.05
actmtch / IPartOpt / tpu / Forward 0.0002233083070022 s 0.0002267879139981 s 0.98
actmtch / DefOpt / tpu / Forward 0.0002221531969989 s 0.0002236571620014 s 0.99
actmtch / IDefOpt / tpu / Forward 0.0002391727619979 s 0.000213723457 s 1.12
actmtch / JaXPipe / tpu / PreRev 0.0002283244559985 s 0.0002418992099992 s 0.94
actmtch / JaXPipe / tpu / PostRev 0.0002339890029979 s 0.0002424740909991 s 0.97
actmtch / JaXPipe / tpu / BothRev 0.0002327623739984 s 0.0002398889989999 s 0.97
actmtch / Jax / tpu / BothRev 0.0002387874509986 s 0.0002405861990009 s 0.99
actmtch / HLOOpt / tpu / PreRev 0.0002357324820004 s 0.0002514944139984 s 0.94
actmtch / HLOOpt / tpu / PostRev 0.0002270937460016 s 0.0002531852960019 s 0.90
actmtch / HLOOpt / tpu / BothRev 0.0002195025389992 s 0.0002331870960006 s 0.94
actmtch / PartOpt / tpu / PreRev 0.0002368842020005 s 0.0002426730909974 s 0.98
actmtch / PartOpt / tpu / PostRev 0.0002083339730015 s 0.0002253151219993 s 0.92
actmtch / PartOpt / tpu / BothRev 0.0002104562120002 s 0.0002278119930015 s 0.92
actmtch / IPartOpt / tpu / PreRev 0.0002270513559997 s 0.0002150392670009 s 1.06
actmtch / IPartOpt / tpu / PostRev 0.0002229987770006 s 0.0002080720949998 s 1.07
actmtch / IPartOpt / tpu / BothRev 0.0002225970670006 s 0.0001973517900005 s 1.13
actmtch / DefOpt / tpu / PreRev 0.000229006245001 s 0.0001956971979998 s 1.17
actmtch / DefOpt / tpu / PostRev 0.0002297174649975 s 0.0001961085490002 s 1.17
actmtch / DefOpt / tpu / BothRev 0.000230071504 s 0.0002008860909991 s 1.15
actmtch / IDefOpt / tpu / PreRev 0.0002217738580002 s 0.0002231891010014 s 0.99
actmtch / IDefOpt / tpu / PostRev 0.0002335337029981 s 0.0001982127900009 s 1.18
actmtch / IDefOpt / tpu / BothRev 0.0002364471219989 s 0.0002016674620026 s 1.17
actmtch / JaXPipe / cpu / Primal 0.0000065499770007591 s 0.000005010517999835429 s 1.31
actmtch / Jax / cpu / Primal 0.000006287145999522181 s 0.000004809449999811477 s 1.31
actmtch / HLOOpt / cpu / Primal 0.000007580161998703261 s 0.00000530116799973257 s 1.43
actmtch / PartOpt / cpu / Primal 0.000006733132000590558 s 0.000004738015999919299 s 1.42
actmtch / IPartOpt / cpu / Primal 0.000006356849000439979 s 0.000004711642000074789 s 1.35
actmtch / DefOpt / cpu / Primal 0.000007675865001147031 s 0.0000053095960001883215 s 1.45
actmtch / IDefOpt / cpu / Primal 0.000007948780999868177 s 0.00000531172199998764 s 1.50
actmtch / JaXPipe / cpu / Forward 0.00001117591899856052 s 0.000008042900999953417 s 1.39
actmtch / Jax / cpu / Forward 0.000009178038999380078 s 0.0000071996839997154895 s 1.27
actmtch / HLOOpt / cpu / Forward 0.000011043126998629304 s 0.000008179603999906248 s 1.35
actmtch / PartOpt / cpu / Forward 0.00001118532899999991 s 0.000008225519999996323 s 1.36
actmtch / IPartOpt / cpu / Forward 0.00001086720300008892 s 0.000008223061000080633 s 1.32
actmtch / DefOpt / cpu / Forward 0.000011542686001121182 s 0.000008229308999943897 s 1.40
actmtch / IDefOpt / cpu / Forward 0.000011569964999580407 s 0.000008161577000009856 s 1.42
actmtch / JaXPipe / cpu / PreRev 0.00001096953300111636 s 0.000008187203000034061 s 1.34
actmtch / JaXPipe / cpu / PostRev 0.000009325298000476324 s 0.000007099333000041952 s 1.31
actmtch / JaXPipe / cpu / BothRev 0.000011693954998918344 s 0.000008179985999959171 s 1.43
actmtch / Jax / cpu / BothRev 0.00000980496000011044 s 0.000007223351999982696 s 1.36
actmtch / HLOOpt / cpu / PreRev 0.000011666439000691752 s 0.00000831829100025061 s 1.40
actmtch / HLOOpt / cpu / PostRev 0.000011098532000687557 s 0.000008230335000007471 s 1.35
actmtch / HLOOpt / cpu / BothRev 0.000011741636999431648 s 0.000008109762999993109 s 1.45
actmtch / PartOpt / cpu / PreRev 0.000011591225998927258 s 0.000008014170000024023 s 1.45
actmtch / PartOpt / cpu / PostRev 0.000009630853000999197 s 0.000007239625000238447 s 1.33
actmtch / PartOpt / cpu / BothRev 0.000011119868000605492 s 0.000008151960999839503 s 1.36
actmtch / IPartOpt / cpu / PreRev 0.00001108679699973436 s 0.000008258112000021356 s 1.34
actmtch / IPartOpt / cpu / PostRev 0.00000914624399956665 s 0.000007345655000335683 s 1.25
actmtch / IPartOpt / cpu / BothRev 0.000011744274999728076 s 0.000008073200000126234 s 1.45
actmtch / DefOpt / cpu / PreRev 0.000011089796000305797 s 0.000008086352000191255 s 1.37
actmtch / DefOpt / cpu / PostRev 0.000011111828000139212 s 0.000008116398999845841 s 1.37
actmtch / DefOpt / cpu / BothRev 0.000011118617998363334 s 0.000008047573000112606 s 1.38
actmtch / IDefOpt / cpu / PreRev 0.000011730315000022527 s 0.000008090865000212944 s 1.45
actmtch / IDefOpt / cpu / PostRev 0.000011017251999874132 s 0.000008463122000193834 s 1.30
actmtch / IDefOpt / cpu / BothRev 0.00001102423500014993 s 0.000008146463999764818 s 1.35
add_one / JaXPipe / cpu / Primal 0.00000913139600015711 s 0.000004943890000049578 s 1.85
add_one / Jax / cpu / Primal 0.000009269508002034856 s 0.000004839874000026612 s 1.92
add_one / HLOOpt / cpu / Primal 0.000009795139001653295 s 0.000004952508000314993 s 1.98
add_one / PartOpt / cpu / Primal 0.000009797143000469076 s 0.000004876157000126114 s 2.01
add_one / IPartOpt / cpu / Primal 0.000009678660997451516 s 0.0000049335029998474055 s 1.96
add_one / DefOpt / cpu / Primal 0.000009727729000587717 s 0.000004943899999943824 s 1.97
add_one / IDefOpt / cpu / Primal 0.000009751150999363744 s 0.000004961550000189164 s 1.97
add_one / JaXPipe / cpu / Forward 0.000014195162999385504 s 0.000008009665999907156 s 1.77
add_one / Jax / cpu / Forward 0.000014371429999300744 s 0.000008026039000014861 s 1.79
add_one / HLOOpt / cpu / Forward 0.000014738262998434949 s 0.000008002075000149489 s 1.84
add_one / PartOpt / cpu / Forward 0.000014944139998988249 s 0.000008069707999766251 s 1.85
add_one / IPartOpt / cpu / Forward 0.000014112556997133652 s 0.000007955976999710401 s 1.77
add_one / DefOpt / cpu / Forward 0.000014263437999034068 s 0.00000793168100017283 s 1.80
add_one / IDefOpt / cpu / Forward 0.00001467586500075413 s 0.000007981792999999015 s 1.84
add_one / JaXPipe / cpu / PreRev 0.000016056928001489723 s 0.000009179855000184032 s 1.75
add_one / JaXPipe / cpu / PostRev 0.000015907626002444887 s 0.000008824837000247498 s 1.80
add_one / JaXPipe / cpu / BothRev 0.000015992617998563218 s 0.000008803057000022819 s 1.82
add_one / Jax / cpu / BothRev 0.000016042252002080204 s 0.000008718068999769457 s 1.84
add_one / HLOOpt / cpu / PreRev 0.000015934178001771214 s 0.000008769959000346716 s 1.82
add_one / HLOOpt / cpu / PostRev 0.00001594909499908681 s 0.000008740447000036511 s 1.82
add_one / HLOOpt / cpu / BothRev 0.000015954492999298964 s 0.000008821307999824058 s 1.81
add_one / PartOpt / cpu / PreRev 0.00001511453799685114 s 0.000008676935000039522 s 1.74
add_one / PartOpt / cpu / PostRev 0.000015878878999501468 s 0.00000884395899993251 s 1.80
add_one / PartOpt / cpu / BothRev 0.000015162203002546447 s 0.000008678790000431036 s 1.75
add_one / IPartOpt / cpu / PreRev 0.000015953769001498586 s 0.00000878293499999927 s 1.82
add_one / IPartOpt / cpu / PostRev 0.00001592499499747646 s 0.000008699948999947083 s 1.83
add_one / IPartOpt / cpu / BothRev 0.00001587855999969179 s 0.000008643516000120143 s 1.84
add_one / DefOpt / cpu / PreRev 0.000015933118000248215 s 0.000008714005999991058 s 1.83
add_one / DefOpt / cpu / PostRev 0.000015867092999542365 s 0.000008628391000002012 s 1.84
add_one / DefOpt / cpu / BothRev 0.00001583232600023621 s 0.000008603973000390396 s 1.84
add_one / IDefOpt / cpu / PreRev 0.000015907223998510745 s 0.000008768942000187963 s 1.81
add_one / IDefOpt / cpu / PostRev 0.000015893154999503167 s 0.000008814397999685752 s 1.80
add_one / IDefOpt / cpu / BothRev 0.000015868701997533208 s 0.000008869111999956658 s 1.79
add_one / JaXPipe / tpu / Primal 0.0001413840069981 s 0.0001347157410018 s 1.05
add_one / Jax / tpu / Primal 0.0001388910389978 s 0.0001308889490028 s 1.06
add_one / HLOOpt / tpu / Primal 0.0001398185379985 s 0.0001315268489997 s 1.06
add_one / PartOpt / tpu / Primal 0.0001420286869979 s 0.0001326260300011 s 1.07
add_one / IPartOpt / tpu / Primal 0.0001379906490001 s 0.0001352304110005 s 1.02
add_one / DefOpt / tpu / Primal 0.000136221030003 s 0.0001320660700002 s 1.03
add_one / IDefOpt / tpu / Primal 0.0001331997209999 s 0.0001362769719999 s 0.98
add_one / JaXPipe / tpu / Forward 0.0001952523370018 s 0.0001970005789989 s 0.99
add_one / Jax / tpu / Forward 0.0002265535559999 s 0.0002013254310004 s 1.13
add_one / HLOOpt / tpu / Forward 0.0001864656910001 s 0.000204026702002 s 0.91
add_one / PartOpt / tpu / Forward 0.0001860581710025 s 0.0002018756819998 s 0.92
add_one / IPartOpt / tpu / Forward 0.0002212863179993 s 0.0001999304609998 s 1.11
add_one / DefOpt / tpu / Forward 0.0002309964739979 s 0.0001976042690002 s 1.17
add_one / IDefOpt / tpu / Forward 0.0002232734370008 s 0.000196255458999 s 1.14
add_one / JaXPipe / tpu / PreRev 0.0002230225669991 s 0.000197708358999 s 1.13
add_one / JaXPipe / tpu / PostRev 0.0002224328170013 s 0.0002060789130009 s 1.08
add_one / JaXPipe / tpu / BothRev 0.0002355142319975 s 0.0002097029550022 s 1.12
add_one / Jax / tpu / BothRev 0.0002153746500007 s 0.0002316825849993 s 0.93
add_one / HLOOpt / tpu / PreRev 0.0002354711029984 s 0.0002017635409974 s 1.17
add_one / HLOOpt / tpu / PostRev 0.0002168967690013 s 0.000225576732002 s 0.96
add_one / HLOOpt / tpu / BothRev 0.000222406877001 s 0.0002305056739969 s 0.96
add_one / PartOpt / tpu / PreRev 0.0002343971440022 s 0.0002291829929999 s 1.02
add_one / PartOpt / tpu / PostRev 0.0002429452299984 s 0.0002225524109999 s 1.09
add_one / PartOpt / tpu / BothRev 0.0002355690719996 s 0.0002242790320015 s 1.05
add_one / IPartOpt / tpu / PreRev 0.0002330739130011 s 0.0002189406100005 s 1.06
add_one / IPartOpt / tpu / PostRev 0.0002224190970009 s 0.0002414977599983 s 0.92
add_one / IPartOpt / tpu / BothRev 0.000240399901002 s 0.0002368421079991 s 1.02
add_one / DefOpt / tpu / PreRev 0.0002169130500005 s 0.0002317573250002 s 0.94
add_one / DefOpt / tpu / PostRev 0.0002210955080008 s 0.000245061100999 s 0.90
add_one / DefOpt / tpu / BothRev 0.000222172927999 s 0.000214042497002 s 1.04
add_one / IDefOpt / tpu / PreRev 0.0002263968960032 s 0.0002234960809983 s 1.01
add_one / IDefOpt / tpu / PostRev 0.0002231096169998 s 0.0002221740700006 s 1.00
add_one / IDefOpt / tpu / BothRev 0.0002342095429994 s 0.0002327537650016 s 1.01
add_one / JaXPipe / cpu / Primal 0.000006645265000770451 s 0.000004943890000049578 s 1.34
add_one / Jax / cpu / Primal 0.000006701545000396436 s 0.000004839874000026612 s 1.38
add_one / HLOOpt / cpu / Primal 0.000006606169999940903 s 0.000004952508000314993 s 1.33
add_one / PartOpt / cpu / Primal 0.00000690369699987059 s 0.000004876157000126114 s 1.42
add_one / IPartOpt / cpu / Primal 0.000006596537001314573 s 0.0000049335029998474055 s 1.34
add_one / DefOpt / cpu / Primal 0.000006646679999903427 s 0.000004943899999943824 s 1.34
add_one / IDefOpt / cpu / Primal 0.0000065689700004440966 s 0.000004961550000189164 s 1.32
add_one / JaXPipe / cpu / Forward 0.000010586246000457324 s 0.000008009665999907156 s 1.32
add_one / Jax / cpu / Forward 0.00001011483099864563 s 0.000008026039000014861 s 1.26
add_one / HLOOpt / cpu / Forward 0.000010664125999028327 s 0.000008002075000149489 s 1.33
add_one / PartOpt / cpu / Forward 0.000010524061999603872 s 0.000008069707999766251 s 1.30
add_one / IPartOpt / cpu / Forward 0.00001006655099990894 s 0.000007955976999710401 s 1.27
add_one / DefOpt / cpu / Forward 0.000010087460999784526 s 0.00000793168100017283 s 1.27
add_one / IDefOpt / cpu / Forward 0.000010138961000848213 s 0.000007981792999999015 s 1.27
add_one / JaXPipe / cpu / PreRev 0.00001132834800046112 s 0.000009179855000184032 s 1.23
add_one / JaXPipe / cpu / PostRev 0.000011246530999414973 s 0.000008824837000247498 s 1.27
add_one / JaXPipe / cpu / BothRev 0.000011405150999053147 s 0.000008803057000022819 s 1.30
add_one / Jax / cpu / BothRev 0.000011366295999323484 s 0.000008718068999769457 s 1.30
add_one / HLOOpt / cpu / PreRev 0.000011345335000441992 s 0.000008769959000346716 s 1.29
add_one / HLOOpt / cpu / PostRev 0.00001137647700124944 s 0.000008740447000036511 s 1.30
add_one / HLOOpt / cpu / BothRev 0.000011377007000191952 s 0.000008821307999824058 s 1.29
add_one / PartOpt / cpu / PreRev 0.00001134534400080156 s 0.000008676935000039522 s 1.31
add_one / PartOpt / cpu / PostRev 0.00001065917899904889 s 0.00000884395899993251 s 1.21
add_one / PartOpt / cpu / BothRev 0.000010714355001255171 s 0.000008678790000431036 s 1.23
add_one / IPartOpt / cpu / PreRev 0.000010687422000046354 s 0.00000878293499999927 s 1.22
add_one / IPartOpt / cpu / PostRev 0.000011298900999463512 s 0.000008699948999947083 s 1.30
add_one / IPartOpt / cpu / BothRev 0.00001135248599894112 s 0.000008643516000120143 s 1.31
add_one / DefOpt / cpu / PreRev 0.000011398364000342553 s 0.000008714005999991058 s 1.31
add_one / DefOpt / cpu / PostRev 0.000011320920999423834 s 0.000008628391000002012 s 1.31
add_one / DefOpt / cpu / BothRev 0.000011322375999952784 s 0.000008603973000390396 s 1.32
add_one / IDefOpt / cpu / PreRev 0.00001222779899944726 s 0.000008768942000187963 s 1.39
add_one / IDefOpt / cpu / PostRev 0.00001120926600015082 s 0.000008814397999685752 s 1.27
add_one / IDefOpt / cpu / BothRev 0.000011366242999429233 s 0.000008869111999956658 s 1.28
add_two / JaXPipe / cpu / Primal 0.00000939140699847485 s 0.000005117662999964523 s 1.84
add_two / Jax / cpu / Primal 0.000009895202001644066 s 0.000005114691000017046 s 1.93
add_two / HLOOpt / cpu / Primal 0.000009345757996925384 s 0.000004998238000098354 s 1.87
add_two / PartOpt / cpu / Primal 0.000009995911001169589 s 0.000005100304999814398 s 1.96
add_two / IPartOpt / cpu / Primal 0.000009996876997320214 s 0.0000050412929999765765 s 1.98
add_two / DefOpt / cpu / Primal 0.000009912302997690858 s 0.0000051184949998059895 s 1.94
add_two / IDefOpt / cpu / Primal 0.000010027586002252064 s 0.000005196600000090257 s 1.93
add_two / JaXPipe / cpu / Forward 0.000014465012998698512 s 0.000008166673999767227 s 1.77
add_two / Jax / cpu / Forward 0.000015016289999039146 s 0.000008187105999695632 s 1.83
add_two / HLOOpt / cpu / Forward 0.000015217415999359218 s 0.00000821973199981585 s 1.85
add_two / PartOpt / cpu / Forward 0.000015114703001017914 s 0.00000820936999980404 s 1.84
add_two / IPartOpt / cpu / Forward 0.000015091963003214916 s 0.000008071210000252904 s 1.87
add_two / DefOpt / cpu / Forward 0.00001497177599958377 s 0.000008141997000166157 s 1.84
add_two / IDefOpt / cpu / Forward 0.00001500459800081444 s 0.000008109810999940237 s 1.85
add_two / JaXPipe / cpu / PreRev 0.000018890925002779112 s 0.000010823311999956786 s 1.75
add_two / JaXPipe / cpu / PostRev 0.000018552251000073737 s 0.000010415172999728385 s 1.78
add_two / JaXPipe / cpu / BothRev 0.00001861418699991191 s 0.000010666382000181329 s 1.75
add_two / Jax / cpu / BothRev 0.000018517332999181238 s 0.00001068545099997209 s 1.73
add_two / HLOOpt / cpu / PreRev 0.00001786987500236137 s 0.00001065027899994675 s 1.68
add_two / HLOOpt / cpu / PostRev 0.0000185308369982522 s 0.000010642900000220834 s 1.74
add_two / HLOOpt / cpu / BothRev 0.00001852431000224897 s 0.000010630987000240566 s 1.74
add_two / PartOpt / cpu / PreRev 0.000018566989001556067 s 0.000010746469999958208 s 1.73
add_two / PartOpt / cpu / PostRev 0.000018644051997398493 s 0.000010739463999925648 s 1.74
add_two / PartOpt / cpu / BothRev 0.000018068101002427285 s 0.000010636098999839305 s 1.70
add_two / IPartOpt / cpu / PreRev 0.00001847051499862573 s 0.00001066233099982128 s 1.73
add_two / IPartOpt / cpu / PostRev 0.00001867189500262612 s 0.000010614927999995416 s 1.76
add_two / IPartOpt / cpu / BothRev 0.000018021903997578196 s 0.000010587897999812411 s 1.70
add_two / DefOpt / cpu / PreRev 0.000018727022001257863 s 0.000010769459000130155 s 1.74
add_two / DefOpt / cpu / PostRev 0.000018607564001285938 s 0.000010686925000300106 s 1.74
add_two / DefOpt / cpu / BothRev 0.000018486144999769748 s 0.000010657176000222535 s 1.73
add_two / IDefOpt / cpu / PreRev 0.00001851738499681233 s 0.00001072462700039978 s 1.73
add_two / IDefOpt / cpu / PostRev 0.000018644755000423176 s 0.00001055853300022136 s 1.77
add_two / IDefOpt / cpu / BothRev 0.00001868870199905359 s 0.000010697150999931184 s 1.75
add_two / JaXPipe / tpu / Primal 0.0001460009059992 s 0.0001336820100004 s 1.09
add_two / Jax / tpu / Primal 0.0001484940650007 s 0.0001336554299996 s 1.11
add_two / HLOOpt / tpu / Primal 0.000153146032997 s 0.0001375829419994 s 1.11
add_two / PartOpt / tpu / Primal 0.0001586184309999 s 0.0001335894599978 s 1.19
add_two / IPartOpt / tpu / Primal 0.0001546265920005 s 0.000134477771 s 1.15
add_two / DefOpt / tpu / Primal 0.0001540349130009 s 0.0001320491100013 s 1.17
add_two / IDefOpt / tpu / Primal 0.0001511163440009 s 0.0001324026500005 s 1.14
add_two / JaXPipe / tpu / Forward 0.0002194482480008 s 0.0002097931749995 s 1.05
add_two / Jax / tpu / Forward 0.0002237698470016 s 0.0002076026239992 s 1.08
add_two / HLOOpt / tpu / Forward 0.0002167735490002 s 0.0002165406679996 s 1.00
add_two / PartOpt / tpu / Forward 0.0002170953099994 s 0.0002189979400027 s 0.99
add_two / IPartOpt / tpu / Forward 0.0002171998389967 s 0.0002203632090022 s 0.99
add_two / DefOpt / tpu / Forward 0.0002157576099998 s 0.0002110960559984 s 1.02
add_two / IDefOpt / tpu / Forward 0.000224572826999 s 0.000213957306998 s 1.05
add_two / JaXPipe / tpu / PreRev 0.0002401943910008 s 0.0002438304909992 s 0.99
add_two / JaXPipe / tpu / PostRev 0.0002404455309988 s 0.0002401589690016 s 1.00
add_two / JaXPipe / tpu / BothRev 0.0002339336030017 s 0.0002414715890008 s 0.97
add_two / Jax / tpu / BothRev 0.000233909393999 s 0.0002399441689995 s 0.97
add_two / HLOOpt / tpu / PreRev 0.0002391410309974 s 0.0002205796099988 s 1.08
add_two / HLOOpt / tpu / PostRev 0.0002405824109991 s 0.0002338687359988 s 1.03
add_two / HLOOpt / tpu / BothRev 0.0002289277050003 s 0.0002416106089985 s 0.95
add_two / PartOpt / tpu / PreRev 0.0002279370150026 s 0.0002399921579999 s 0.95
add_two / PartOpt / tpu / PostRev 0.0002383869710029 s 0.0002315413639989 s 1.03
add_two / PartOpt / tpu / BothRev 0.0002367696220026 s 0.0002419188189996 s 0.98
add_two / IPartOpt / tpu / PreRev 0.0002339918929974 s 0.0002394905690016 s 0.98
add_two / IPartOpt / tpu / PostRev 0.0002275381459985 s 0.0002289005429993 s 0.99
add_two / IPartOpt / tpu / BothRev 0.0002531258460003 s 0.0002324576559985 s 1.09
add_two / DefOpt / tpu / PreRev 0.0002514812869994 s 0.0002303956850009 s 1.09
add_two / DefOpt / tpu / PostRev 0.0002367466719988 s 0.0002412197500016 s 0.98
add_two / DefOpt / tpu / BothRev 0.0002484807380023 s 0.0002424211899997 s 1.02
add_two / IDefOpt / tpu / PreRev 0.0002282787860021 s 0.0002508631929995 s 0.91
add_two / IDefOpt / tpu / PostRev 0.0002298499049975 s 0.0002398078879996 s 0.96
add_two / IDefOpt / tpu / BothRev 0.0002301615750002 s 0.0002301049740017 s 1.00
add_two / JaXPipe / cpu / Primal 0.000006914449000760215 s 0.000005117662999964523 s 1.35
add_two / Jax / cpu / Primal 0.000007215084999188548 s 0.000005114691000017046 s 1.41
add_two / HLOOpt / cpu / Primal 0.000007197421999080689 s 0.000004998238000098354 s 1.44
add_two / PartOpt / cpu / Primal 0.0000068408610004553336 s 0.000005100304999814398 s 1.34
add_two / IPartOpt / cpu / Primal 0.000006862070000352105 s 0.0000050412929999765765 s 1.36
add_two / DefOpt / cpu / Primal 0.000006845964000603999 s 0.0000051184949998059895 s 1.34
add_two / IDefOpt / cpu / Primal 0.000006863782999062096 s 0.000005196600000090257 s 1.32
add_two / JaXPipe / cpu / Forward 0.00001023166099912487 s 0.000008166673999767227 s 1.25
add_two / Jax / cpu / Forward 0.000010673263001081069 s 0.000008187105999695632 s 1.30
add_two / HLOOpt / cpu / Forward 0.000010688880000088827 s 0.00000821973199981585 s 1.30
add_two / PartOpt / cpu / Forward 0.000010322485999495256 s 0.00000820936999980404 s 1.26
add_two / IPartOpt / cpu / Forward 0.00001072506499986048 s 0.000008071210000252904 s 1.33
add_two / DefOpt / cpu / Forward 0.000010217374001513237 s 0.000008141997000166157 s 1.25
add_two / IDefOpt / cpu / Forward 0.000010209557000052882 s 0.000008109810999940237 s 1.26
add_two / JaXPipe / cpu / PreRev 0.000012814146999517108 s 0.000010823311999956786 s 1.18
add_two / JaXPipe / cpu / PostRev 0.00001333831800002372 s 0.000010415172999728385 s 1.28
add_two / JaXPipe / cpu / BothRev 0.000013114363000568118 s 0.000010666382000181329 s 1.23
add_two / Jax / cpu / BothRev 0.000013199847000578302 s 0.00001068545099997209 s 1.24
add_two / HLOOpt / cpu / PreRev 0.00001320185500117077 s 0.00001065027899994675 s 1.24
add_two / HLOOpt / cpu / PostRev 0.000012787129999196622 s 0.000010642900000220834 s 1.20
add_two / HLOOpt / cpu / BothRev 0.000013246751001133816 s 0.000010630987000240566 s 1.25
add_two / PartOpt / cpu / PreRev 0.0000133030629986024 s 0.000010746469999958208 s 1.24
add_two / PartOpt / cpu / PostRev 0.000013257308999527597 s 0.000010739463999925648 s 1.23
add_two / PartOpt / cpu / BothRev 0.000012572564999572933 s 0.000010636098999839305 s 1.18
add_two / IPartOpt / cpu / PreRev 0.00001324731100066856 s 0.00001066233099982128 s 1.24
add_two / IPartOpt / cpu / PostRev 0.000013145270999302738 s 0.000010614927999995416 s 1.24
add_two / IPartOpt / cpu / BothRev 0.000013124478000463569 s 0.000010587897999812411 s 1.24
add_two / DefOpt / cpu / PreRev 0.000012613602999408612 s 0.000010769459000130155 s 1.17
add_two / DefOpt / cpu / PostRev 0.000013338716000362185 s 0.000010686925000300106 s 1.25
add_two / DefOpt / cpu / BothRev 0.00001328589699915028 s 0.000010657176000222535 s 1.25
add_two / IDefOpt / cpu / PreRev 0.000013230263999503225 s 0.00001072462700039978 s 1.23
add_two / IDefOpt / cpu / PostRev 0.000013466016000165835 s 0.00001055853300022136 s 1.28
add_two / IDefOpt / cpu / BothRev 0.00001271478599846887 s 0.000010697150999931184 s 1.19
cache / JaXPipe / cpu / Primal 0.000009379070997965754 s 0.000004610288000094443 s 2.03
cache / Jax / cpu / Primal 0.000009308889999374516 s 0.0000048870710002120174 s 1.90
cache / HLOOpt / cpu / Primal 0.000009459836001042276 s 0.000004565947000173765 s 2.07
cache / PartOpt / cpu / Primal 0.000010068804000184172 s 0.000004967840000063006 s 2.03
cache / IPartOpt / cpu / Primal 0.000009955909001291727 s 0.000004909473000225262 s 2.03
cache / DefOpt / cpu / Primal 0.000009476893999817548 s 0.000004601310999987618 s 2.06
cache / IDefOpt / cpu / Primal 0.000009505654001259244 s 0.00000455357400005596 s 2.09
cache / JaXPipe / cpu / Forward 0.000015705288002209273 s 0.000010998106999977608 s 1.43
cache / Jax / cpu / Forward 0.00001539076200060663 s 0.000010775844000363577 s 1.43
cache / HLOOpt / cpu / Forward 0.000014756350999959977 s 0.00001119033600025432 s 1.32
cache / PartOpt / cpu / Forward 0.000015224288999888813 s 0.000011269240999808971 s 1.35
cache / IPartOpt / cpu / Forward 0.000014988879000156885 s 0.00001113953499998388 s 1.35
cache / DefOpt / cpu / Forward 0.000015245164999214469 s 0.000011068823000186968 s 1.38
cache / IDefOpt / cpu / Forward 0.00001547584899890353 s 0.000010735372999988613 s 1.44
cache / JaXPipe / cpu / PreRev 0.00001584566699966672 s 0.000011910075999821856 s 1.33
cache / JaXPipe / cpu / PostRev 0.000017390811000950635 s 0.00001574259400013034 s 1.10
cache / JaXPipe / cpu / BothRev 0.00001529660600135685 s 0.000012121721999847068 s 1.26
cache / Jax / cpu / BothRev 0.000017374011000356404 s 0.00001538118200005556 s 1.13
cache / HLOOpt / cpu / PreRev 0.000015706845999375218 s 0.000012285113999951137 s 1.28
cache / HLOOpt / cpu / PostRev 0.000015844830999412806 s 0.000011640605000138748 s 1.36
cache / HLOOpt / cpu / BothRev 0.00001513257799888379 s 0.00001212775200019678 s 1.25
cache / PartOpt / cpu / PreRev 0.00001539905599929625 s 0.000012027923999994528 s 1.28
cache / PartOpt / cpu / PostRev 0.000017808255997806554 s 0.000015316942000026757 s 1.16
cache / PartOpt / cpu / BothRev 0.000015436568999575683 s 0.000011914537000393466 s 1.30
cache / IPartOpt / cpu / PreRev 0.000015854478999244746 s 0.000011896418000105768 s 1.33
cache / IPartOpt / cpu / PostRev 0.000018067331999191083 s 0.000015347735999966973 s 1.18
cache / IPartOpt / cpu / BothRev 0.000016026870001951464 s 0.000011832433000108722 s 1.35
cache / DefOpt / cpu / PreRev 0.000015602481998939765 s 0.000012308014000154798 s 1.27
cache / DefOpt / cpu / PostRev 0.000016122791999805486 s 0.000012036612999963837 s 1.34
cache / DefOpt / cpu / BothRev 0.000015860240000620252 s 0.000012262581999948452 s 1.29
cache / IDefOpt / cpu / PreRev 0.00001512919999731821 s 0.000011921959000119386 s 1.27
cache / IDefOpt / cpu / PostRev 0.00001599743900078465 s 0.0000117301079999379 s 1.36
cache / IDefOpt / cpu / BothRev 0.00001575717900050222 s 0.000012103028999717936 s 1.30
cache / JaXPipe / tpu / Primal 0.0001367856889992 s 0.0001376403419999 s 0.99
cache / Jax / tpu / Primal 0.0001382475589998 s 0.0001351508110019 s 1.02
cache / HLOOpt / tpu / Primal 0.0001393926379969 s 0.0001327002299985 s 1.05
cache / PartOpt / tpu / Primal 0.0001400170980014 s 0.0001383640519998 s 1.01
cache / IPartOpt / tpu / Primal 0.0001387625990028 s 0.0001371617619988 s 1.01
cache / DefOpt / tpu / Primal 0.0001386435389977 s 0.000137978693001 s 1.00
cache / IDefOpt / tpu / Primal 0.0001384561890008 s 0.0001361918719994 s 1.02
cache / JaXPipe / tpu / Forward 0.000198687367003 s 0.0002234952910002 s 0.89
cache / Jax / tpu / Forward 0.0002028839140002 s 0.0002449497509987 s 0.83
cache / HLOOpt / tpu / Forward 0.0002171655189995 s 0.0002435148000004 s 0.89
cache / PartOpt / tpu / Forward 0.0002092319820003 s 0.0002447079899975 s 0.86
cache / IPartOpt / tpu / Forward 0.0002020652150022 s 0.0002396994379996 s 0.84
cache / DefOpt / tpu / Forward 0.0002216492980005 s 0.0002273752029977 s 0.97
cache / IDefOpt / tpu / Forward 0.0002090383030008 s 0.0002273260229994 s 0.92
cache / JaXPipe / tpu / PreRev 0.0002102027820001 s 0.0002229471110003 s 0.94
cache / JaXPipe / tpu / PostRev 0.0002103978720006 s 0.0002202414100029 s 0.96
cache / JaXPipe / tpu / BothRev 0.0002220742370009 s 0.0002280470429977 s 0.97
cache / Jax / tpu / BothRev 0.0002297910149973 s 0.0002226944810026 s 1.03
cache / HLOOpt / tpu / PreRev 0.0002461110190015 s 0.0002150539780013 s 1.14
cache / HLOOpt / tpu / PostRev 0.0002227755370004 s 0.0002002519810012 s 1.11
cache / HLOOpt / tpu / BothRev 0.0002221396869972 s 0.0002038520720016 s 1.09
cache / PartOpt / tpu / PreRev 0.0002190804679994 s 0.0001986616300018 s 1.10
cache / PartOpt / tpu / PostRev 0.0002094894319998 s 0.0001972245489996 s 1.06
cache / PartOpt / tpu / BothRev 0.0002111902719989 s 0.0002354446669996 s 0.90
cache / IPartOpt / tpu / PreRev 0.0002103642709989 s 0.0002398993989991 s 0.88
cache / IPartOpt / tpu / PostRev 0.0002218861480032 s 0.0002011875609969 s 1.10
cache / IPartOpt / tpu / BothRev 0.0002212434280008 s 0.0002170584080013 s 1.02
cache / DefOpt / tpu / PreRev 0.0002192847389997 s 0.0002197246100004 s 1.00
cache / DefOpt / tpu / PostRev 0.0002217457280021 s 0.0002500064630003 s 0.89
cache / DefOpt / tpu / BothRev 0.0002248723160009 s 0.0002190466690008 s 1.03
cache / IDefOpt / tpu / PreRev 0.0002232017170026 s 0.0002174275690013 s 1.03
cache / IDefOpt / tpu / PostRev 0.0002210981370008 s 0.000220117140001 s 1.00
cache / IDefOpt / tpu / BothRev 0.0002284985050027 s 0.0002251049219994 s 1.02
cache / JaXPipe / cpu / Primal 0.000006201958000019658 s 0.000004610288000094443 s 1.35
cache / Jax / cpu / Primal 0.000006357264001053409 s 0.0000048870710002120174 s 1.30
cache / HLOOpt / cpu / Primal 0.000006269757999689318 s 0.000004565947000173765 s 1.37
cache / PartOpt / cpu / Primal 0.000006245497999771033 s 0.000004967840000063006 s 1.26
cache / IPartOpt / cpu / Primal 0.000006740855998941697 s 0.000004909473000225262 s 1.37
cache / DefOpt / cpu / Primal 0.000006710436999128433 s 0.000004601310999987618 s 1.46
cache / IDefOpt / cpu / Primal 0.000006282471000304213 s 0.00000455357400005596 s 1.38
cache / JaXPipe / cpu / Forward 0.000010602156000459217 s 0.000010998106999977608 s 0.96
cache / Jax / cpu / Forward 0.000010662988001058693 s 0.000010775844000363577 s 0.99
cache / HLOOpt / cpu / Forward 0.000010698492000301484 s 0.00001119033600025432 s 0.96
cache / PartOpt / cpu / Forward 0.00001035479900019709 s 0.000011269240999808971 s 0.92
cache / IPartOpt / cpu / Forward 0.000010910032999163378 s 0.00001113953499998388 s 0.98
cache / DefOpt / cpu / Forward 0.000010373355000410813 s 0.000011068823000186968 s 0.94
cache / IDefOpt / cpu / Forward 0.000011273538999375887 s 0.000010735372999988613 s 1.05
cache / JaXPipe / cpu / PreRev 0.000011345492999680571 s 0.000011910075999821856 s 0.95
cache / JaXPipe / cpu / PostRev 0.00001283505700121168 s 0.00001574259400013034 s 0.82
cache / JaXPipe / cpu / BothRev 0.000011564619000637324 s 0.000012121721999847068 s 0.95
cache / Jax / cpu / BothRev 0.000012405854000462568 s 0.00001538118200005556 s 0.81
cache / HLOOpt / cpu / PreRev 0.000010767432000648116 s 0.000012285113999951137 s 0.88
cache / HLOOpt / cpu / PostRev 0.000010654252000676934 s 0.000011640605000138748 s 0.92
cache / HLOOpt / cpu / BothRev 0.000011014902000169968 s 0.00001212775200019678 s 0.91
cache / PartOpt / cpu / PreRev 0.00001121335999960138 s 0.000012027923999994528 s 0.93
cache / PartOpt / cpu / PostRev 0.000012331443000221045 s 0.000015316942000026757 s 0.81
cache / PartOpt / cpu / BothRev 0.000010940060999928392 s 0.000011914537000393466 s 0.92
cache / IPartOpt / cpu / PreRev 0.000010698129000957125 s 0.000011896418000105768 s 0.90
cache / IPartOpt / cpu / PostRev 0.00001255916099944443 s 0.000015347735999966973 s 0.82
cache / IPartOpt / cpu / BothRev 0.000011094875999333452 s 0.000011832433000108722 s 0.94
cache / DefOpt / cpu / PreRev 0.000010976891000609612 s 0.000012308014000154798 s 0.89
cache / DefOpt / cpu / PostRev 0.000010958847000438254 s 0.000012036612999963837 s 0.91
cache / DefOpt / cpu / BothRev 0.00001053699099975347 s 0.000012262581999948452 s 0.86
cache / IDefOpt / cpu / PreRev 0.000010684794000553666 s 0.000011921959000119386 s 0.90
cache / IDefOpt / cpu / PostRev 0.000010966265999741154 s 0.0000117301079999379 s 0.93
cache / IDefOpt / cpu / BothRev 0.000011236985001232824 s 0.000012103028999717936 s 0.93
Concat / JaXPipe / cpu / Primal 0.000009184627000649926 s 0.000004935824000313005 s 1.86
Concat / Jax / cpu / Primal 0.000009221217998856446 s 0.000004809287000171025 s 1.92
Concat / HLOOpt / cpu / Primal 0.000009695001001091442 s 0.0000049511939996591535 s 1.96
Concat / PartOpt / cpu / Primal 0.00000922022299710079 s 0.000004963275000136491 s 1.86
Concat / IPartOpt / cpu / Primal 0.000009211585998855298 s 0.000004876702000274235 s 1.89
Concat / DefOpt / cpu / Primal 0.00000941018299999996 s 0.000004906742000002851 s 1.92
Concat / IDefOpt / cpu / Primal 0.000009733447001053718 s 0.000004970080000020971 s 1.96
Concat / JaXPipe / cpu / Forward 0.000014735837001353504 s 0.000008068141999956423 s 1.83
Concat / Jax / cpu / Forward 0.00001418587700027274 s 0.000008060895000198798 s 1.76
Concat / HLOOpt / cpu / Forward 0.00001414692000253126 s 0.000007997551999778807 s 1.77
Concat / PartOpt / cpu / Forward 0.000014159917998767923 s 0.000007910280000032799 s 1.79
Concat / IPartOpt / cpu / Forward 0.000014079476000915748 s 0.000008070426999893244 s 1.74
Concat / DefOpt / cpu / Forward 0.000014742038001713808 s 0.00000827997599981245 s 1.78
Concat / IDefOpt / cpu / Forward 0.000015016592002211838 s 0.000007918111999970278 s 1.90
Concat / JaXPipe / cpu / PreRev 0.000015220968998619356 s 0.000008632422000118823 s 1.76
Concat / JaXPipe / cpu / PostRev 0.000016028082998673197 s 0.000008634448999600863 s 1.86
Concat / JaXPipe / cpu / BothRev 0.00001601896099964506 s 0.00000871425300010742 s 1.84
Concat / Jax / cpu / BothRev 0.00001585008499750984 s 0.000008734413999718526 s 1.81
Concat / HLOOpt / cpu / PreRev 0.000015947130999848015 s 0.000008913596000184043 s 1.79
Concat / HLOOpt / cpu / PostRev 0.00001598433099934482 s 0.000008801600999959192 s 1.82
Concat / HLOOpt / cpu / BothRev 0.0000159307110006921 s 0.000008755009000196879 s 1.82
Concat / PartOpt / cpu / PreRev 0.00001593057799982489 s 0.000008634915999664371 s 1.84
Concat / PartOpt / cpu / PostRev 0.00001604275999852689 s 0.00000867884299987054 s 1.85
Concat / PartOpt / cpu / BothRev 0.000015382207999209642 s 0.000008710005000011733 s 1.77
Concat / IPartOpt / cpu / PreRev 0.000016078707001724977 s 0.000008878565000031812 s 1.81
Concat / IPartOpt / cpu / PostRev 0.000015222295998682966 s 0.000008714056999906461 s 1.75
Concat / IPartOpt / cpu / BothRev 0.000016109943000628845 s 0.00000884776500015505 s 1.82
Concat / DefOpt / cpu / PreRev 0.00001628281699959189 s 0.000008693026999935683 s 1.87
Concat / DefOpt / cpu / PostRev 0.000015949985001498134 s 0.000008715844999642286 s 1.83
Concat / DefOpt / cpu / BothRev 0.000015254311001626777 s 0.000008616912999968918 s 1.77
Concat / IDefOpt / cpu / PreRev 0.00001592868000079761 s 0.00000886408999986088 s 1.80
Concat / IDefOpt / cpu / PostRev 0.00001598420300069847 s 0.000008752216000175395 s 1.83
Concat / IDefOpt / cpu / BothRev 0.00001598729299803381 s 0.0000086947700001474 s 1.84
Concat / JaXPipe / tpu / Primal 0.0001511420429997 s 0.0001407409740022 s 1.07
Concat / Jax / tpu / Primal 0.0001479920450001 s 0.0001396091829992 s 1.06
Concat / HLOOpt / tpu / Primal 0.000145204706001 s 0.0001401479029991 s 1.04
Concat / PartOpt / tpu / Primal 0.0001453385860004 s 0.0001412327340003 s 1.03
Concat / IPartOpt / tpu / Primal 0.0001474958149992 s 0.0001407218040003 s 1.05
Concat / DefOpt / tpu / Primal 0.0001351516899994 s 0.0001613575629999 s 0.84
Concat / IDefOpt / tpu / Primal 0.0001347274500003 s 0.0001570831710014 s 0.86
Concat / JaXPipe / tpu / Forward 0.0002092622130003 s 0.0002196864789984 s 0.95
Concat / Jax / tpu / Forward 0.0002032477450011 s 0.0002062139940026 s 0.99
Concat / HLOOpt / tpu / Forward 0.000205712414001 s 0.0002287297139992 s 0.90
Concat / PartOpt / tpu / Forward 0.0002096034120004 s 0.0002434736399991 s 0.86
Concat / IPartOpt / tpu / Forward 0.000238021131001 s 0.0002022011709996 s 1.18
Concat / DefOpt / tpu / Forward 0.0002264413550001 s 0.0002032638320015 s 1.11
Concat / IDefOpt / tpu / Forward 0.0002262958550018 s 0.0002034978720002 s 1.11
Concat / JaXPipe / tpu / PreRev 0.000231995803002 s 0.0002209748299974 s 1.05
Concat / JaXPipe / tpu / PostRev 0.0002285250259992 s 0.0002260619619992 s 1.01
Concat / JaXPipe / tpu / BothRev 0.0002281029849982 s 0.000223421781 s 1.02
Concat / Jax / tpu / BothRev 0.0002170287090011 s 0.0002250291719974 s 0.96
Concat / HLOOpt / tpu / PreRev 0.0002162858600022 s 0.0002439791209981 s 0.89
Concat / HLOOpt / tpu / PostRev 0.0002154942199995 s 0.0002241178809999 s 0.96
Concat / HLOOpt / tpu / BothRev 0.0002166093100022 s 0.0002243803320016 s 0.97
Concat / PartOpt / tpu / PreRev 0.0002189683280012 s 0.0002240992919978 s 0.98
Concat / PartOpt / tpu / PostRev 0.0002168494289981 s 0.0002295271139992 s 0.94
Concat / PartOpt / tpu / BothRev 0.0002166308199994 s 0.0002236507419984 s 0.97
Concat / IPartOpt / tpu / PreRev 0.000198446716 s 0.0002218605599991 s 0.89
Concat / IPartOpt / tpu / PostRev 0.0001978588569982 s 0.0002498973030014 s 0.79
Concat / IPartOpt / tpu / BothRev 0.0002478064210008 s 0.0002242405520009 s 1.11
Concat / DefOpt / tpu / PreRev 0.0002309572880003 s 0.0002230337910004 s 1.04
Concat / DefOpt / tpu / PostRev 0.0002318645069972 s 0.0002499784629981 s 0.93
Concat / DefOpt / tpu / BothRev 0.0002297371979984 s 0.0002352164260009 s 0.98
Concat / IDefOpt / tpu / PreRev 0.0002279004589981 s 0.0002362897370003 s 0.96
Concat / IDefOpt / tpu / PostRev 0.0002380274949973 s 0.0002155183270006 s 1.10
Concat / IDefOpt / tpu / BothRev 0.0002361033259985 s 0.0002193996000023 s 1.08
Concat / JaXPipe / cpu / Primal 0.000006420286999855307 s 0.000004935824000313005 s 1.30
Concat / Jax / cpu / Primal 0.000006541438999192905 s 0.000004809287000171025 s 1.36
Concat / HLOOpt / cpu / Primal 0.000006566502001078333 s 0.0000049511939996591535 s 1.33
Concat / PartOpt / cpu / Primal 0.000006943084999875282 s 0.000004963275000136491 s 1.40
Concat / IPartOpt / cpu / Primal 0.000006881197999973665 s 0.000004876702000274235 s 1.41
Concat / DefOpt / cpu / Primal 0.000006412614999135258 s 0.000004906742000002851 s 1.31
Concat / IDefOpt / cpu / Primal 0.000006718188000377268 s 0.000004970080000020971 s 1.35
Concat / JaXPipe / cpu / Forward 0.000010422110999570578 s 0.000008068141999956423 s 1.29
Concat / Jax / cpu / Forward 0.000010497425999346888 s 0.000008060895000198798 s 1.30
Concat / HLOOpt / cpu / Forward 0.00000996592400042573 s 0.000007997551999778807 s 1.25
Concat / PartOpt / cpu / Forward 0.000010079367999423992 s 0.000007910280000032799 s 1.27
Concat / IPartOpt / cpu / Forward 0.000009885211999062447 s 0.000008070426999893244 s 1.22
Concat / DefOpt / cpu / Forward 0.000010695042999941508 s 0.00000827997599981245 s 1.29
Concat / IDefOpt / cpu / Forward 0.00001038659999903757 s 0.000007918111999970278 s 1.31
Concat / JaXPipe / cpu / PreRev 0.000010596123000141234 s 0.000008632422000118823 s 1.23
Concat / JaXPipe / cpu / PostRev 0.000011248140999668977 s 0.000008634448999600863 s 1.30
Concat / JaXPipe / cpu / BothRev 0.000011291200000414392 s 0.00000871425300010742 s 1.30
Concat / Jax / cpu / BothRev 0.000011450477000835236 s 0.000008734413999718526 s 1.31
Concat / HLOOpt / cpu / PreRev 0.000011488050000480143 s 0.000008913596000184043 s 1.29
Concat / HLOOpt / cpu / PostRev 0.000011394150998967234 s 0.000008801600999959192 s 1.29
Concat / HLOOpt / cpu / BothRev 0.000011540843999682691 s 0.000008755009000196879 s 1.32
Concat / PartOpt / cpu / PreRev 0.000011231036998651689 s 0.000008634915999664371 s 1.30
Concat / PartOpt / cpu / PostRev 0.000011432138999225572 s 0.00000867884299987054 s 1.32
Concat / PartOpt / cpu / BothRev 0.000010631316999933915 s 0.000008710005000011733 s 1.22
Concat / IPartOpt / cpu / PreRev 0.00001123716899928695 s 0.000008878565000031812 s 1.27
Concat / IPartOpt / cpu / PostRev 0.000011300756999844452 s 0.000008714056999906461 s 1.30
Concat / IPartOpt / cpu / BothRev 0.000011409712000386208 s 0.00000884776500015505 s 1.29
Concat / DefOpt / cpu / PreRev 0.000011287305998848753 s 0.000008693026999935683 s 1.30
Concat / DefOpt / cpu / PostRev 0.000011310854000839757 s 0.000008715844999642286 s 1.30
Concat / DefOpt / cpu / BothRev 0.000011265658999036531 s 0.000008616912999968918 s 1.31
Concat / IDefOpt / cpu / PreRev 0.000011438874000305076 s 0.00000886408999986088 s 1.29
Concat / IDefOpt / cpu / PostRev 0.000011452851998910774 s 0.000008752216000175395 s 1.31
Concat / IDefOpt / cpu / BothRev 0.000011304510000627488 s 0.0000086947700001474 s 1.30
const_scatter / JaXPipe / cpu / Primal 0.000008617358998890268 s 0.00000457273299980443 s 1.88
const_scatter / Jax / cpu / Primal 0.00000891187299930607 s 0.000004540475999874616 s 1.96
const_scatter / HLOOpt / cpu / Primal 0.000010020575999078572 s 0.000005096807999962039 s 1.97
const_scatter / PartOpt / cpu / Primal 0.000009443240000109654 s 0.0000045606260000568 s 2.07
const_scatter / IPartOpt / cpu / Primal 0.000008938471000874416 s 0.000004527223999957641 s 1.97
const_scatter / DefOpt / cpu / Primal 0.00001004152099994826 s 0.000005392006999954902 s 1.86
const_scatter / IDefOpt / cpu / Primal 0.000010558926001976942 s 0.000005162890999599767 s 2.05
const_scatter / JaXPipe / cpu / Forward 0.00001569367900083307 s 0.000008455307000076573 s 1.86
const_scatter / Jax / cpu / Forward 0.000013460880996717608 s 0.000007473365999885572 s 1.80
const_scatter / HLOOpt / cpu / Forward 0.00001489468200088595 s 0.000008526113999778317 s 1.75
const_scatter / PartOpt / cpu / Forward 0.000015482350001548186 s 0.000008523288000105822 s 1.82
const_scatter / IPartOpt / cpu / Forward 0.000015530235999904107 s 0.000008113974000025338 s 1.91
const_scatter / DefOpt / cpu / Forward 0.000015647208001610125 s 0.000008190042000023823 s 1.91
const_scatter / IDefOpt / cpu / Forward 0.000015665560000343247 s 0.000008501315000103205 s 1.84
const_scatter / JaXPipe / tpu / Primal 0.0001500389800021 s 0.000141045094002 s 1.06
const_scatter / Jax / tpu / Primal 0.0001335628569977 s 0.0001536014289995 s 0.87
const_scatter / HLOOpt / tpu / Primal 0.0001329380669994 s 0.0001539307790008 s 0.86
const_scatter / PartOpt / tpu / Primal 0.0001370631259997 s 0.0001417008640019 s 0.97
const_scatter / IPartOpt / tpu / Primal 0.0001411056840006 s 0.0001413195940003 s 1.00
const_scatter / DefOpt / tpu / Primal 0.0001473834219978 s 0.0001518743889973 s 0.97
const_scatter / IDefOpt / tpu / Primal 0.0001392684940001 s 0.0001530235000027 s 0.91
const_scatter / JaXPipe / tpu / Forward 0.0002336223569982 s 0.0002473130020007 s 0.94
const_scatter / Jax / tpu / Forward 0.0002220883920017 s 0.0002267710630003 s 0.98
const_scatter / HLOOpt / tpu / Forward 0.0002239902310029 s 0.000226067202002 s 0.99
const_scatter / PartOpt / tpu / Forward 0.0002207120509992 s 0.0002261523930028 s 0.98
const_scatter / IPartOpt / tpu / Forward 0.0002322948069995 s 0.0002263781830006 s 1.03
const_scatter / DefOpt / tpu / Forward 0.0002389079139975 s 0.0002319068139986 s 1.03
const_scatter / IDefOpt / tpu / Forward 0.0002311083169988 s 0.0002312087850004 s 1.00
const_scatter / JaXPipe / cpu / Primal 0.0000066218970005138544 s 0.00000457273299980443 s 1.45
const_scatter / Jax / cpu / Primal 0.000006194272998982342 s 0.000004540475999874616 s 1.36
const_scatter / HLOOpt / cpu / Primal 0.000007071522000842379 s 0.000005096807999962039 s 1.39
const_scatter / PartOpt / cpu / Primal 0.000006215010000232724 s 0.0000045606260000568 s 1.36
const_scatter / IPartOpt / cpu / Primal 0.000006262196999159642 s 0.000004527223999957641 s 1.38
const_scatter / DefOpt / cpu / Primal 0.000007049589999951422 s 0.000005392006999954902 s 1.31
const_scatter / IDefOpt / cpu / Primal 0.000007147893000365002 s 0.000005162890999599767 s 1.38
const_scatter / JaXPipe / cpu / Forward 0.000010764155000288157 s 0.000008455307000076573 s 1.27
const_scatter / Jax / cpu / Forward 0.000009482767000008607 s 0.000007473365999885572 s 1.27
const_scatter / HLOOpt / cpu / Forward 0.000010333999000067709 s 0.000008526113999778317 s 1.21
const_scatter / PartOpt / cpu / Forward 0.000010886125000979518 s 0.000008523288000105822 s 1.28
const_scatter / IPartOpt / cpu / Forward 0.000010396968000350172 s 0.000008113974000025338 s 1.28
const_scatter / DefOpt / cpu / Forward 0.000010351372000513948 s 0.000008190042000023823 s 1.26
const_scatter / IDefOpt / cpu / Forward 0.00001032235899947409 s 0.000008501315000103205 s 1.21
GenDot / JaXPipe / cpu / Primal 0.000009735113999340682 s 0.000005363540999951511 s 1.82
GenDot / Jax / cpu / Primal 0.000009774287998880029 s 0.000004987264000192226 s 1.96
GenDot / HLOOpt / cpu / Primal 0.000010791929002152756 s 0.000005337630000212812 s 2.02
GenDot / PartOpt / cpu / Primal 0.00001033085100061726 s 0.000005115057000239176 s 2.02
GenDot / IPartOpt / cpu / Primal 0.000009849933001532918 s 0.000005099805000099877 s 1.93
GenDot / DefOpt / cpu / Primal 0.000011224302998016356 s 0.000005303803000060725 s 2.12
GenDot / IDefOpt / cpu / Primal 0.000011755404000723502 s 0.0000053544789998341 s 2.20
GenDot / JaXPipe / cpu / Forward 0.000016898426998523064 s 0.000008204633999866928 s 2.06
GenDot / Jax / cpu / Forward 0.000014472366001427871 s 0.000007319183999697998 s 1.98
GenDot / HLOOpt / cpu / Forward 0.00001678034200085676 s 0.00000819195599979139 s 2.05
GenDot / PartOpt / cpu / Forward 0.000016742053001507883 s 0.000008356307999747515 s 2.00
GenDot / IPartOpt / cpu / Forward 0.000016900751001230673 s 0.000008089055999789707 s 2.09
GenDot / DefOpt / cpu / Forward 0.000016407733000960433 s 0.0000081760299999587 s 2.01
GenDot / IDefOpt / cpu / Forward 0.000016590179002378135 s 0.000008373077000214835 s 1.98
GenDot / JaXPipe / cpu / PreRev 0.000016094161997898483 s 0.000008202785999856132 s 1.96
GenDot / JaXPipe / cpu / PostRev 0.000014530213000398362 s 0.0000073416229997747 s 1.98
GenDot / JaXPipe / cpu / BothRev 0.00001679136600068887 s 0.000008172309999736172 s 2.05
GenDot / Jax / cpu / BothRev 0.000014633055001468165 s 0.000007649227000001701 s 1.91
GenDot / HLOOpt / cpu / PreRev 0.00001607587200123817 s 0.000008203029000014795 s 1.96
GenDot / HLOOpt / cpu / PostRev 0.0000159936270028993 s 0.000008208765999825119 s 1.95
GenDot / HLOOpt / cpu / BothRev 0.000016001895000954392 s 0.000008150720999765327 s 1.96
GenDot / PartOpt / cpu / PreRev 0.000016881413001101464 s 0.0000080856610002229 s 2.09
GenDot / PartOpt / cpu / PostRev 0.000014777061998756836 s 0.000007493225999951391 s 1.97
GenDot / PartOpt / cpu / BothRev 0.000016291820000333247 s 0.000008152233000146225 s 2.00
GenDot / IPartOpt / cpu / PreRev 0.000015966644998115952 s 0.000008036902999720041 s 1.99
GenDot / IPartOpt / cpu / PostRev 0.000014668995998363245 s 0.000007573905999834096 s 1.94
GenDot / IPartOpt / cpu / BothRev 0.00001677431600182899 s 0.000008123784999952477 s 2.06
GenDot / DefOpt / cpu / PreRev 0.000016919583998969757 s 0.000008181425999737258 s 2.07
GenDot / DefOpt / cpu / PostRev 0.000016034199001296655 s 0.000008134196999890264 s 1.97
GenDot / DefOpt / cpu / BothRev 0.000016849075000209267 s 0.00000825376700004199 s 2.04
GenDot / IDefOpt / cpu / PreRev 0.000016780862999439704 s 0.000008039405000090483 s 2.09
GenDot / IDefOpt / cpu / PostRev 0.000016857219001394698 s 0.000008142678999774943 s 2.07
GenDot / IDefOpt / cpu / BothRev 0.000016741501000069546 s 0.000008186646000012842 s 2.04
GenDot / JaXPipe / tpu / Primal 0.0001510908899981 s 0.0001385719530007 s 1.09
GenDot / Jax / tpu / Primal 0.0001343022969995 s 0.0001370201319987 s 0.98
GenDot / HLOOpt / tpu / Primal 0.0001349477660005 s 0.0001468499869988 s 0.92
GenDot / PartOpt / tpu / Primal 0.000135856724999 s 0.0001472271870006 s 0.92
GenDot / IPartOpt / tpu / Primal 0.0001338539959979 s 0.0001487187769998 s 0.90
GenDot / DefOpt / tpu / Primal 0.0001340798559976 s 0.0001370135719989 s 0.98
GenDot / IDefOpt / tpu / Primal 0.0001396611339987 s 0.0001355829409985 s 1.03
GenDot / JaXPipe / tpu / Forward 0.0002030087290004 s 0.0002209969399991 s 0.92
GenDot / Jax / tpu / Forward 0.0002041653580017 s 0.0002250514709994 s 0.91
GenDot / HLOOpt / tpu / Forward 0.0002038642290026 s 0.0002159735179993 s 0.94
GenDot / PartOpt / tpu / Forward 0.0002026136589993 s 0.0002164614180001 s 0.94
GenDot / IPartOpt / tpu / Forward 0.0002147687640026 s 0.0002241911109995 s 0.96
GenDot / DefOpt / tpu / Forward 0.0002022226389999 s 0.0002248555719997 s 0.90
GenDot / IDefOpt / tpu / Forward 0.0002048474090006 s 0.0002309675139986 s 0.89
GenDot / JaXPipe / tpu / PreRev 0.0002265405690013 s 0.0001990916899994 s 1.14
GenDot / JaXPipe / tpu / PostRev 0.0002339358769968 s 0.000202191201999 s 1.16
GenDot / JaXPipe / tpu / BothRev 0.0002349736260002 s 0.0002002699799995 s 1.17
GenDot / Jax / tpu / BothRev 0.0002302475680007 s 0.0002042072020012 s 1.13
GenDot / HLOOpt / tpu / PreRev 0.000230860587002 s 0.0002025805319972 s 1.14
GenDot / HLOOpt / tpu / PostRev 0.0002304296580005 s 0.0002124322759991 s 1.08
GenDot / HLOOpt / tpu / BothRev 0.0002362012560006 s 0.0002023397320008 s 1.17
GenDot / PartOpt / tpu / PreRev 0.0002102521260021 s 0.0002007194010002 s 1.05
GenDot / PartOpt / tpu / PostRev 0.0002037562389996 s 0.0002016342210008 s 1.01
GenDot / PartOpt / tpu / BothRev 0.0002023050200004 s 0.0002185399289992 s 0.93
GenDot / IPartOpt / tpu / PreRev 0.0002036551979981 s 0.0002206987099998 s 0.92
GenDot / IPartOpt / tpu / PostRev 0.0002115639160001 s 0.0002323288460029 s 0.91
GenDot / IPartOpt / tpu / BothRev 0.0002085488669981 s 0.0002344413459977 s 0.89
GenDot / DefOpt / tpu / PreRev 0.0001968551109966 s 0.0002331723250026 s 0.84
GenDot / DefOpt / tpu / PostRev 0.0002208713110012 s 0.000231323134998 s 0.95
GenDot / DefOpt / tpu / BothRev 0.0002221923510005 s 0.0002216223409996 s 1.00
GenDot / IDefOpt / tpu / PreRev 0.0002021766190009 s 0.000197928810001 s 1.02
GenDot / IDefOpt / tpu / PostRev 0.0002202732930018 s 0.0001978303889991 s 1.11
GenDot / IDefOpt / tpu / BothRev 0.0002226667909999 s 0.000215751058 s 1.03
GenDot / JaXPipe / cpu / Primal 0.000006690862001050846 s 0.000005363540999951511 s 1.25
GenDot / Jax / cpu / Primal 0.000007209251998574473 s 0.000004987264000192226 s 1.45
GenDot / HLOOpt / cpu / Primal 0.000008017451998966863 s 0.000005337630000212812 s 1.50
GenDot / PartOpt / cpu / Primal 0.000006934517999980016 s 0.000005115057000239176 s 1.36
GenDot / IPartOpt / cpu / Primal 0.00000687536999976146 s 0.000005099805000099877 s 1.35
GenDot / DefOpt / cpu / Primal 0.000007763596000586403 s 0.000005303803000060725 s 1.46
GenDot / IDefOpt / cpu / Primal 0.000007654065000679111 s 0.0000053544789998341 s 1.43
GenDot / JaXPipe / cpu / Forward 0.000011205314000108048 s 0.000008204633999866928 s 1.37
GenDot / Jax / cpu / Forward 0.000010102448999532498 s 0.000007319183999697998 s 1.38
GenDot / HLOOpt / cpu / Forward 0.000011585006999666802 s 0.00000819195599979139 s 1.41
GenDot / PartOpt / cpu / Forward 0.000011029520999727536 s 0.000008356307999747515 s 1.32
GenDot / IPartOpt / cpu / Forward 0.000011680164001518278 s 0.000008089055999789707 s 1.44
GenDot / DefOpt / cpu / Forward 0.000011246771999140037 s 0.0000081760299999587 s 1.38
GenDot / IDefOpt / cpu / Forward 0.000011429805999796372 s 0.000008373077000214835 s 1.37
GenDot / JaXPipe / cpu / PreRev 0.000011213873000087914 s 0.000008202785999856132 s 1.37
GenDot / JaXPipe / cpu / PostRev 0.00001023152299967478 s 0.0000073416229997747 s 1.39
GenDot / JaXPipe / cpu / BothRev 0.000011685486000715172 s 0.000008172309999736172 s 1.43
GenDot / Jax / cpu / BothRev 0.000010172655000133091 s 0.000007649227000001701 s 1.33
GenDot / HLOOpt / cpu / PreRev 0.000011063275000196882 s 0.000008203029000014795 s 1.35
GenDot / HLOOpt / cpu / PostRev 0.000011642228999335202 s 0.000008208765999825119 s 1.42
GenDot / HLOOpt / cpu / BothRev 0.000011015440999472048 s 0.000008150720999765327 s 1.35
GenDot / PartOpt / cpu / PreRev 0.000011715855998772895 s 0.0000080856610002229 s 1.45
GenDot / PartOpt / cpu / PostRev 0.00001015668600120989 s 0.000007493225999951391 s 1.36
GenDot / PartOpt / cpu / BothRev 0.000011074976000600146 s 0.000008152233000146225 s 1.36
GenDot / IPartOpt / cpu / PreRev 0.000011117833000753308 s 0.000008036902999720041 s 1.38
GenDot / IPartOpt / cpu / PostRev 0.00001030497700048727 s 0.000007573905999834096 s 1.36
GenDot / IPartOpt / cpu / BothRev 0.000011654545000055806 s 0.000008123784999952477 s 1.43
GenDot / DefOpt / cpu / PreRev 0.000011624921000475298 s 0.000008181425999737258 s 1.42
GenDot / DefOpt / cpu / PostRev 0.000011104029999842168 s 0.000008134196999890264 s 1.37
GenDot / DefOpt / cpu / BothRev 0.000011630664001131663 s 0.00000825376700004199 s 1.41
GenDot / IDefOpt / cpu / PreRev 0.000011677348000375788 s 0.000008039405000090483 s 1.45
GenDot / IDefOpt / cpu / PostRev 0.000011767797999709727 s 0.000008142678999774943 s 1.45
GenDot / IDefOpt / cpu / BothRev 0.000011547960000825695 s 0.000008186646000012842 s 1.41
hlo_ffi / JaXPipe / cpu / Primal 0.000013719972004764714 s 0.000007839142000193533 s 1.75
hlo_ffi / Jax / cpu / Primal 0.00001345660000515636 s 0.000007418492999931914 s 1.81
hlo_ffi / HLOOpt / cpu / Primal 0.00001428442400356289 s 0.000007519155999943905 s 1.90
hlo_ffi / PartOpt / cpu / Primal 0.000014317298999230845 s 0.000007463871999789262 s 1.92
hlo_ffi / IPartOpt / cpu / Primal 0.000013579619000665844 s 0.000007491344999834835 s 1.81
hlo_ffi / DefOpt / cpu / Primal 0.0000143187339999713 s 0.0000075683639997805584 s 1.89
hlo_ffi / IDefOpt / cpu / Primal 0.000013575972996477505 s 0.000007518132000313927 s 1.81
hlo_ffi / JaXPipe / cpu / Forward 0.000019585028996516485 s 0.000011934550999740168 s 1.64
hlo_ffi / Jax / cpu / Forward 0.000020403542002895847 s 0.000011807264999788456 s 1.73
hlo_ffi / HLOOpt / cpu / Forward 0.00001943866699730279 s 0.000011862851999921986 s 1.64
hlo_ffi / PartOpt / cpu / Forward 0.00002036422200035304 s 0.00001188141300008283 s 1.71
hlo_ffi / IPartOpt / cpu / Forward 0.000019435990005149507 s 0.000011791235000146115 s 1.65
hlo_ffi / DefOpt / cpu / Forward 0.000020422231995326 s 0.0000118316110001615 s 1.73
hlo_ffi / IDefOpt / cpu / Forward 0.000019455636000202504 s 0.000011888229999840403 s 1.64
hlo_ffi / JaXPipe / cpu / PreRev 0.000018888735001382883 s 0.000011295965000044816 s 1.67
hlo_ffi / JaXPipe / cpu / PostRev 0.000018868183004087768 s 0.000011389642000267486 s 1.66
hlo_ffi / JaXPipe / cpu / BothRev 0.00001982461900479393 s 0.000011269571999946493 s 1.76
hlo_ffi / Jax / cpu / BothRev 0.00002001159099745564 s 0.000011273670000264249 s 1.78
hlo_ffi / HLOOpt / cpu / PreRev 0.00001979750399914337 s 0.000011427941999954785 s 1.73
hlo_ffi / HLOOpt / cpu / PostRev 0.00001973249799630139 s 0.00001137640300021303 s 1.73
hlo_ffi / HLOOpt / cpu / BothRev 0.00001973871899826918 s 0.000011270746999798576 s 1.75
hlo_ffi / PartOpt / cpu / PreRev 0.000019731168998987413 s 0.00001122469199981424 s 1.76
hlo_ffi / PartOpt / cpu / PostRev 0.00001986446299997624 s 0.000011275952000232792 s 1.76
hlo_ffi / PartOpt / cpu / BothRev 0.0000189084319936228 s 0.000011229502999867691 s 1.68
hlo_ffi / IPartOpt / cpu / PreRev 0.000019858387000567743 s 0.000011358218000168565 s 1.75
hlo_ffi / IPartOpt / cpu / PostRev 0.00001970732799964025 s 0.00001142774299978555 s 1.72
hlo_ffi / IPartOpt / cpu / BothRev 0.00001970170099957613 s 0.000011210428000140382 s 1.76
hlo_ffi / DefOpt / cpu / PreRev 0.000019835931998386512 s 0.000011252631999923324 s 1.76
hlo_ffi / DefOpt / cpu / PostRev 0.0000198772639996605 s 0.000011269901000105164 s 1.76
hlo_ffi / DefOpt / cpu / BothRev 0.0000199474429973634 s 0.000011277249999693595 s 1.77
hlo_ffi / IDefOpt / cpu / PreRev 0.00002011094599583885 s 0.000011280108999926596 s 1.78
hlo_ffi / IDefOpt / cpu / PostRev 0.00001982053599931533 s 0.000011384377000013044 s 1.74
hlo_ffi / IDefOpt / cpu / BothRev 0.00001969208700029412 s 0.000011318458000005194 s 1.74
hlo_ffi / JaXPipe / tpu / Primal 0.000213233587001 s 0.0002321582549993 s 0.92
hlo_ffi / Jax / tpu / Primal 0.0002084949580021 s 0.0002150232480016 s 0.97
hlo_ffi / HLOOpt / tpu / Primal 0.0002134651960004 s 0.0002239209420004 s 0.95
hlo_ffi / PartOpt / tpu / Primal 0.0002097703279978 s 0.0001948999989981 s 1.08
hlo_ffi / IPartOpt / tpu / Primal 0.0002036274700003 s 0.0001972840899979 s 1.03
hlo_ffi / DefOpt / tpu / Primal 0.0001997975120029 s 0.0001975565699976 s 1.01
hlo_ffi / IDefOpt / tpu / Primal 0.0001998816119994 s 0.0001981083190003 s 1.01
hlo_ffi / JaXPipe / tpu / Forward 0.0002219451730015 s 0.0002308534139992 s 0.96
hlo_ffi / Jax / tpu / Forward 0.0002287294900015 s 0.0002314064250022 s 0.99
hlo_ffi / HLOOpt / tpu / Forward 0.0002280355809998 s 0.000230832874 s 0.99
hlo_ffi / PartOpt / tpu / Forward 0.0002282196710002 s 0.0002255239719997 s 1.01
hlo_ffi / IPartOpt / tpu / Forward 0.0002276524599983 s 0.000223478980999 s 1.02
hlo_ffi / DefOpt / tpu / Forward 0.0002227385030018 s 0.0002261563419997 s 0.98
hlo_ffi / IDefOpt / tpu / Forward 0.000224597362001 s 0.0002388576580015 s 0.94
hlo_ffi / JaXPipe / tpu / PreRev 0.0002228249129984 s 0.0002413110490015 s 0.92
hlo_ffi / JaXPipe / tpu / PostRev 0.0002237469020001 s 0.0002405449590005 s 0.93
hlo_ffi / JaXPipe / tpu / BothRev 0.0002233238119988 s 0.0002269628530011 s 0.98
hlo_ffi / Jax / tpu / BothRev 0.0002453970939968 s 0.0002258397929981 s 1.09
hlo_ffi / HLOOpt / tpu / PreRev 0.0002388232470002 s 0.0002290256939995 s 1.04
hlo_ffi / HLOOpt / tpu / PostRev 0.0002238129830002 s 0.0002344345560013 s 0.95
hlo_ffi / HLOOpt / tpu / BothRev 0.0002461312830018 s 0.0002351431770002 s 1.05
hlo_ffi / PartOpt / tpu / PreRev 0.0002248625419997 s 0.0002530076350012 s 0.89
hlo_ffi / PartOpt / tpu / PostRev 0.0002264348710014 s 0.000255136345997 s 0.89
hlo_ffi / PartOpt / tpu / BothRev 0.0002106329180023 s 0.0002527574140003 s 0.83
hlo_ffi / IPartOpt / tpu / PreRev 0.0002077088790028 s 0.000252035974001 s 0.82
hlo_ffi / IPartOpt / tpu / PostRev 0.0002264454409996 s 0.0002147790369999 s 1.05
hlo_ffi / IPartOpt / tpu / BothRev 0.0002175708239992 s 0.0002259637319984 s 0.96
hlo_ffi / DefOpt / tpu / PreRev 0.0002334325679985 s 0.0002378407469986 s 0.98
hlo_ffi / DefOpt / tpu / PostRev 0.0001884303959996 s 0.0002445132909997 s 0.77
hlo_ffi / DefOpt / tpu / BothRev 0.0001947717939983 s 0.0002460119610004 s 0.79
hlo_ffi / IDefOpt / tpu / PreRev 0.0001915318850005 s 0.0002406988889997 s 0.80
hlo_ffi / IDefOpt / tpu / PostRev 0.0002024363009986 s 0.0002402098989987 s 0.84
hlo_ffi / IDefOpt / tpu / BothRev 0.0001925263349985 s 0.0002630663689997 s 0.73
hlo_ffi / JaXPipe / cpu / Primal 0.000009713323999676504 s 0.000007839142000193533 s 1.24
hlo_ffi / Jax / cpu / Primal 0.000009285272000852272 s 0.000007418492999931914 s 1.25
hlo_ffi / HLOOpt / cpu / Primal 0.00000972137600001588 s 0.000007519155999943905 s 1.29
hlo_ffi / PartOpt / cpu / Primal 0.000009709199999633713 s 0.000007463871999789262 s 1.30
hlo_ffi / IPartOpt / cpu / Primal 0.000009706224000183284 s 0.000007491344999834835 s 1.30
hlo_ffi / DefOpt / cpu / Primal 0.0000097438200009492 s 0.0000075683639997805584 s 1.29
hlo_ffi / IDefOpt / cpu / Primal 0.000009759032000147272 s 0.000007518132000313927 s 1.30
hlo_ffi / JaXPipe / cpu / Forward 0.000013700796000193805 s 0.000011934550999740168 s 1.15
hlo_ffi / Jax / cpu / Forward 0.000013504861999535934 s 0.000011807264999788456 s 1.14
hlo_ffi / HLOOpt / cpu / Forward 0.000013596385999335326 s 0.000011862851999921986 s 1.15
hlo_ffi / PartOpt / cpu / Forward 0.000014324931998999092 s 0.00001188141300008283 s 1.21
hlo_ffi / IPartOpt / cpu / Forward 0.000014259302000937169 s 0.000011791235000146115 s 1.21
hlo_ffi / DefOpt / cpu / Forward 0.00001417684899934102 s 0.0000118316110001615 s 1.20
hlo_ffi / IDefOpt / cpu / Forward 0.000014281797000876396 s 0.000011888229999840403 s 1.20
hlo_ffi / JaXPipe / cpu / PreRev 0.000013270908999402308 s 0.000011295965000044816 s 1.17
hlo_ffi / JaXPipe / cpu / PostRev 0.000014080680999541073 s 0.000011389642000267486 s 1.24
hlo_ffi / JaXPipe / cpu / BothRev 0.000013940268998339888 s 0.000011269571999946493 s 1.24
hlo_ffi / Jax / cpu / BothRev 0.000013967926999612246 s 0.000011273670000264249 s 1.24
hlo_ffi / HLOOpt / cpu / PreRev 0.000013155224000001907 s 0.000011427941999954785 s 1.15
hlo_ffi / HLOOpt / cpu / PostRev 0.000013780348999716808 s 0.00001137640300021303 s 1.21
hlo_ffi / HLOOpt / cpu / BothRev 0.000013904677998652917 s 0.000011270746999798576 s 1.23
hlo_ffi / PartOpt / cpu / PreRev 0.000014009157001055429 s 0.00001122469199981424 s 1.25
hlo_ffi / PartOpt / cpu / PostRev 0.000013953375000710368 s 0.000011275952000232792 s 1.24
hlo_ffi / PartOpt / cpu / BothRev 0.000013376331000472418 s 0.000011229502999867691 s 1.19
hlo_ffi / IPartOpt / cpu / PreRev 0.00001383376499870792 s 0.000011358218000168565 s 1.22
hlo_ffi / IPartOpt / cpu / PostRev 0.000013908195000112757 s 0.00001142774299978555 s 1.22
hlo_ffi / IPartOpt / cpu / BothRev 0.000013899327999752132 s 0.000011210428000140382 s 1.24
hlo_ffi / DefOpt / cpu / PreRev 0.000013305616999787164 s 0.000011252631999923324 s 1.18
hlo_ffi / DefOpt / cpu / PostRev 0.000013278214000820298 s 0.000011269901000105164 s 1.18
hlo_ffi / DefOpt / cpu / BothRev 0.000013909587998568896 s 0.000011277249999693595 s 1.23
hlo_ffi / IDefOpt / cpu / PreRev 0.000013795787001072312 s 0.000011280108999926596 s 1.22
hlo_ffi / IDefOpt / cpu / PostRev 0.000013993660000778618 s 0.000011384377000013044 s 1.23
hlo_ffi / IDefOpt / cpu / BothRev 0.000013772233000054256 s 0.000011318458000005194 s 1.22
jaxmd40 / JaXPipe / cpu / Primal 0.0807375165997655 s 0.0776407775992993 s 1.04
jaxmd40 / Jax / cpu / Primal 0.0756915271995239 s 0.0772177848004503 s 0.98
jaxmd40 / HLOOpt / cpu / Primal 0.1120746845990652 s 0.1125256151979556 s 1.00
jaxmd40 / PartOpt / cpu / Primal 0.0790631256008055 s 0.0801443256001221 s 0.99
jaxmd40 / IPartOpt / cpu / Primal 0.0788310837990138 s 0.0748769390018424 s 1.05
jaxmd40 / DefOpt / cpu / Primal 0.1131765688012819 s 0.1125798382010543 s 1.01
jaxmd40 / IDefOpt / cpu / Primal 0.1161975718001485 s 0.1125397459982195 s 1.03
jaxmd40 / JaXPipe / cpu / Forward 0.2056594184003188 s 0.2003848549997201 s 1.03
jaxmd40 / Jax / cpu / Forward 0.1118746619991725 s 0.10732561420009 s 1.04
jaxmd40 / HLOOpt / cpu / Forward 0.2023327652001171 s 0.2016559996001888 s 1.00
jaxmd40 / PartOpt / cpu / Forward 0.2052690295997308 s 0.1943234499980462 s 1.06
jaxmd40 / IPartOpt / cpu / Forward 0.2032871296003577 s 0.1973254209995502 s 1.03
jaxmd40 / DefOpt / cpu / Forward 0.2026857158009079 s 0.1976926774019375 s 1.03
jaxmd40 / IDefOpt / cpu / Forward 0.2038306065995129 s 0.1950129371980438 s 1.05
jaxmd40 / JaXPipe / cpu / PostRev 0.1613540676000411 s 0.1623857514001429 s 0.99
jaxmd40 / Jax / cpu / BothRev 0.1611945213997387 s 0.1563392763986485 s 1.03
jaxmd40 / HLOOpt / cpu / PostRev 0.2196597449990804 s 0.2210803747992031 s 0.99
jaxmd40 / PartOpt / cpu / PostRev 0.152520007599378 s 0.1462311036011669 s 1.04
jaxmd40 / IPartOpt / cpu / PostRev 0.1540635030003613 s 0.1580070208001416 s 0.98
jaxmd40 / DefOpt / cpu / PostRev 0.220637022399751 s 0.2229618484008824 s 0.99
jaxmd40 / IDefOpt / cpu / PostRev 0.2177232942005503 s 0.2224910117976833 s 0.98
jaxmd40 / JaXPipe / tpu / Primal 0.009318164200522 s 0.0093028182003763 s 1.00
jaxmd40 / Jax / tpu / Primal 0.0093000583998218 s 0.0092769041999417 s 1.00
jaxmd40 / HLOOpt / tpu / Primal 0.0092335704001015 s 0.0092318661998433 s 1.00
jaxmd40 / PartOpt / tpu / Primal 0.0092686444004357 s 0.0092618102004053 s 1.00
jaxmd40 / IPartOpt / tpu / Primal 0.0092975441999442 s 0.0092679842004145 s 1.00
jaxmd40 / DefOpt / tpu / Primal 0.0091164423996815 s 0.0091002942004706 s 1.00
jaxmd40 / IDefOpt / tpu / Primal 0.0090615244000218 s 0.0090521022000757 s 1.00
jaxmd40 / JaXPipe / tpu / Forward 0.0178756087996589 s 0.0178711722001025 s 1.00
jaxmd40 / Jax / tpu / Forward 0.0183268745997338 s 0.018313716200646 s 1.00
jaxmd40 / HLOOpt / tpu / Forward 0.0178494407999096 s 0.0178584099994623 s 1.00
jaxmd40 / PartOpt / tpu / Forward 0.0179101688001537 s 0.0178887861999101 s 1.00
jaxmd40 / IPartOpt / tpu / Forward 0.0178632369999832 s 0.0178844021997065 s 1.00
jaxmd40 / DefOpt / tpu / Forward 0.0178726029997051 s 0.0178899541999271 s 1.00
jaxmd40 / IDefOpt / tpu / Forward 0.017857699200249 s 0.0178565839996736 s 1.00
jaxmd40 / JaXPipe / tpu / PostRev 0.0198039123999478 s 0.0197915087999717 s 1.00
jaxmd40 / Jax / tpu / BothRev 0.0197814683997421 s 0.0197698049996688 s 1.00
jaxmd40 / HLOOpt / tpu / PostRev 0.0188859525995212 s 0.0189108624006621 s 1.00
jaxmd40 / PartOpt / tpu / PostRev 0.0197445584002707 s 0.0197206490003736 s 1.00
jaxmd40 / IPartOpt / tpu / PostRev 0.019753108399891 s 0.0197423890000209 s 1.00
jaxmd40 / DefOpt / tpu / PostRev 0.0183148229996731 s 0.0182931622002797 s 1.00
jaxmd40 / IDefOpt / tpu / PostRev 0.018070780999551 s 0.0180457821996242 s 1.00
jaxmd40 / JaXPipe / cpu / Primal 0.0650513130000035 s 0.0776407775992993 s 0.84
jaxmd40 / Jax / cpu / Primal 0.0666043012002774 s 0.0772177848004503 s 0.86
jaxmd40 / HLOOpt / cpu / Primal 0.0887709364000329 s 0.1125256151979556 s 0.79
jaxmd40 / PartOpt / cpu / Primal 0.0654293596002389 s 0.0801443256001221 s 0.82
jaxmd40 / IPartOpt / cpu / Primal 0.0686978189998626 s 0.0748769390018424 s 0.92
jaxmd40 / DefOpt / cpu / Primal 0.093090965599913 s 0.1125798382010543 s 0.83
jaxmd40 / IDefOpt / cpu / Primal 0.0942038504002994 s 0.1125397459982195 s 0.84
jaxmd40 / JaXPipe / cpu / Forward 0.1650755465998372 s 0.2003848549997201 s 0.82
jaxmd40 / Jax / cpu / Forward 0.0984727901999576 s 0.10732561420009 s 0.92
jaxmd40 / HLOOpt / cpu / Forward 0.1614486668000609 s 0.2016559996001888 s 0.80
jaxmd40 / PartOpt / cpu / Forward 0.1651423982002597 s 0.1943234499980462 s 0.85
jaxmd40 / IPartOpt / cpu / Forward 0.1643059263999021 s 0.1973254209995502 s 0.83
jaxmd40 / DefOpt / cpu / Forward 0.1619208493997575 s 0.1976926774019375 s 0.82
jaxmd40 / IDefOpt / cpu / Forward 0.1646060050003143 s 0.1950129371980438 s 0.84
jaxmd40 / JaXPipe / cpu / PostRev 0.1470017165996978 s 0.1623857514001429 s 0.91
jaxmd40 / Jax / cpu / BothRev 0.1400143157999991 s 0.1563392763986485 s 0.90
jaxmd40 / HLOOpt / cpu / PostRev 0.181446041199888 s 0.2210803747992031 s 0.82
jaxmd40 / PartOpt / cpu / PostRev 0.138664100599999 s 0.1462311036011669 s 0.95
jaxmd40 / IPartOpt / cpu / PostRev 0.1398881170000095 s 0.1580070208001416 s 0.89
jaxmd40 / DefOpt / cpu / PostRev 0.1791314932001114 s 0.2229618484008824 s 0.80
jaxmd40 / IDefOpt / cpu / PostRev 0.1768737998001597 s 0.2224910117976833 s 0.79
neural_gcm_dynamic_forcing_deterministic_1_4_deg / JaXPipe / cpu / Primal 51.9003105970005 s 52.04495415899146 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / Jax / cpu / Primal 52.021632390999 s 51.077285592997214 s 1.02
neural_gcm_dynamic_forcing_deterministic_1_4_deg / HLOOpt / cpu / Primal 50.85886086800019 s 50.5468015219958 s 1.01
neural_gcm_dynamic_forcing_deterministic_1_4_deg / PartOpt / cpu / Primal 50.76168627399966 s 49.33023756300099 s 1.03
neural_gcm_dynamic_forcing_deterministic_1_4_deg / IPartOpt / cpu / Primal 50.67958549899777 s 49.68710263600224 s 1.02
neural_gcm_dynamic_forcing_deterministic_1_4_deg / DefOpt / cpu / Primal 25.78109640399998 s 25.64614806800091 s 1.01
neural_gcm_dynamic_forcing_deterministic_1_4_deg / IDefOpt / cpu / Primal 56.12322033900273 s 55.09455818599963 s 1.02
neural_gcm_dynamic_forcing_deterministic_1_4_deg / JaXPipe / tpu / Primal 0.1900293459984823 s 0.1900886369985528 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / Jax / tpu / Primal 0.1899680959977558 s 0.1900179860022035 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / HLOOpt / tpu / Primal 0.1886303070023132 s 0.1887619459994312 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / PartOpt / tpu / Primal 0.2028411950013833 s 0.2028582230013853 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / IPartOpt / tpu / Primal 0.2027527549980732 s 0.2028795330006687 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / DefOpt / tpu / Primal 0.1738992050013621 s 0.1739699390018358 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / IDefOpt / tpu / Primal 0.1848898620010004 s 0.1851474740033154 s 1.00
neural_gcm_dynamic_forcing_deterministic_1_4_deg / JaXPipe / cpu / Primal 52.45201178399839 s 52.04495415899146 s 1.01
neural_gcm_dynamic_forcing_deterministic_1_4_deg / Jax / cpu / Primal 53.02243184200052 s 51.077285592997214 s 1.04
neural_gcm_dynamic_forcing_deterministic_1_4_deg / HLOOpt / cpu / Primal 51.25032147499951 s 50.5468015219958 s 1.01
neural_gcm_dynamic_forcing_deterministic_1_4_deg / PartOpt / cpu / Primal 51.85063355400053 s 49.33023756300099 s 1.05
neural_gcm_dynamic_forcing_deterministic_1_4_deg / IPartOpt / cpu / Primal 51.29915524200078 s 49.68710263600224 s 1.03
neural_gcm_dynamic_forcing_deterministic_1_4_deg / DefOpt / cpu / Primal 23.47013539099862 s 25.64614806800091 s 0.92
neural_gcm_dynamic_forcing_deterministic_1_4_deg / IDefOpt / cpu / Primal 56.12952389599923 s 55.09455818599963 s 1.02
scatter_sum / JaXPipe / cpu / Primal 0.000010908273001405178 s 0.000005582883999977639 s 1.95
scatter_sum / Jax / cpu / Primal 0.00001081518600039999 s 0.000005537862999972276 s 1.95
scatter_sum / HLOOpt / cpu / Primal 0.000010969716000545304 s 0.000005586341999787692 s 1.96
scatter_sum / PartOpt / cpu / Primal 0.000010947915001452202 s 0.00000549940799965043 s 1.99
scatter_sum / IPartOpt / cpu / Primal 0.00001066289099981077 s 0.000005470600000080594 s 1.95
scatter_sum / DefOpt / cpu / Primal 0.00001126267200015718 s 0.000005552988999625086 s 2.03
scatter_sum / IDefOpt / cpu / Primal 0.000010881453999900258 s 0.000005579821999617707 s 1.95
scatter_sum / JaXPipe / cpu / Forward 0.000015569262999633793 s 0.000008504458000061277 s 1.83
scatter_sum / Jax / cpu / Forward 0.00001636581500133616 s 0.000008450866999737627 s 1.94
scatter_sum / HLOOpt / cpu / Forward 0.000015534976999333595 s 0.00000851859600015814 s 1.82
scatter_sum / PartOpt / cpu / Forward 0.000016337007000402082 s 0.000008492328000102133 s 1.92
scatter_sum / IPartOpt / cpu / Forward 0.000016169807000551372 s 0.000008530288999736512 s 1.90
scatter_sum / DefOpt / cpu / Forward 0.00001557574600155931 s 0.000008528337999905488 s 1.83
scatter_sum / IDefOpt / cpu / Forward 0.000016285912999592257 s 0.00000847586500003672 s 1.92
scatter_sum / JaXPipe / cpu / PreRev 0.000016453866999654565 s 0.000008587461999923108 s 1.92
scatter_sum / JaXPipe / cpu / PostRev 0.000016604684999038 s 0.000008512100999723771 s 1.95
scatter_sum / JaXPipe / cpu / BothRev 0.000016542698998819105 s 0.000008555274000173086 s 1.93
scatter_sum / Jax / cpu / BothRev 0.000016640761001326608 s 0.000008547964000172215 s 1.95
scatter_sum / HLOOpt / cpu / PreRev 0.00001575586599938106 s 0.000008617295000021841 s 1.83
scatter_sum / HLOOpt / cpu / PostRev 0.000016621466998913094 s 0.000008582896999996592 s 1.94
scatter_sum / HLOOpt / cpu / BothRev 0.000015810171000339325 s 0.000008562118000099872 s 1.85
scatter_sum / PartOpt / cpu / PreRev 0.00001652171199748409 s 0.000008642949999739358 s 1.91
scatter_sum / PartOpt / cpu / PostRev 0.000016506514999491627 s 0.000008591775999775564 s 1.92
scatter_sum / PartOpt / cpu / BothRev 0.000015801470002770657 s 0.000008602850999977817 s 1.84
scatter_sum / IPartOpt / cpu / PreRev 0.000015826094000658484 s 0.000008546456999738439 s 1.85
scatter_sum / IPartOpt / cpu / PostRev 0.000015793543003383092 s 0.000008552622000024712 s 1.85
scatter_sum / IPartOpt / cpu / BothRev 0.00001660843200079398 s 0.000008500688999902194 s 1.95
scatter_sum / DefOpt / cpu / PreRev 0.000016498507997312117 s 0.000008667650000006688 s 1.90
scatter_sum / DefOpt / cpu / PostRev 0.00001662606500030961 s 0.000008575887000006332 s 1.94
scatter_sum / DefOpt / cpu / BothRev 0.00001671788900057436 s 0.000008634672999960457 s 1.94
scatter_sum / IDefOpt / cpu / PreRev 0.00001656404900131747 s 0.000008544630999949731 s 1.94
scatter_sum / IDefOpt / cpu / PostRev 0.000016472275998239638 s 0.000009006696000142256 s 1.83
scatter_sum / IDefOpt / cpu / BothRev 0.00001656231100059813 s 0.000008609931000137294 s 1.92
scatter_sum / JaXPipe / tpu / Primal 0.000154235957998 s 0.0001396392429996 s 1.10
scatter_sum / Jax / tpu / Primal 0.0001537772979972 s 0.0001364576819978 s 1.13
scatter_sum / HLOOpt / tpu / Primal 0.0001511792500023 s 0.0001482936070024 s 1.02
scatter_sum / PartOpt / tpu / Primal 0.0001489152899994 s 0.0001488284570004 s 1.00
scatter_sum / IPartOpt / tpu / Primal 0.0001359852259993 s 0.0001502765680015 s 0.90
scatter_sum / DefOpt / tpu / Primal 0.0001420445629992 s 0.0001629007539995 s 0.87
scatter_sum / IDefOpt / tpu / Primal 0.0001421058029991 s 0.0001377135029979 s 1.03
scatter_sum / JaXPipe / tpu / Forward 0.0001982534400012 s 0.0002089332449977 s 0.95
scatter_sum / Jax / tpu / Forward 0.0001983484010015 s 0.0002247547410006 s 0.88
scatter_sum / HLOOpt / tpu / Forward 0.0002079395270011 s 0.0002371164779979 s 0.88
scatter_sum / PartOpt / tpu / Forward 0.0002227118320006 s 0.0002041647030018 s 1.09
scatter_sum / IPartOpt / tpu / Forward 0.0002238427099982 s 0.0002159681180019 s 1.04
scatter_sum / DefOpt / tpu / Forward 0.0002056232979994 s 0.000225274311997 s 0.91
scatter_sum / IDefOpt / tpu / Forward 0.0002228776010015 s 0.0002444878009991 s 0.91
scatter_sum / JaXPipe / tpu / PreRev 0.0002190758529977 s 0.0002087017140001 s 1.05
scatter_sum / JaXPipe / tpu / PostRev 0.0002030509779979 s 0.0002228407010006 s 0.91
scatter_sum / JaXPipe / tpu / BothRev 0.0002200361920004 s 0.0002501337629983 s 0.88
scatter_sum / Jax / tpu / BothRev 0.0002251928000005 s 0.0002505022439981 s 0.90
scatter_sum / HLOOpt / tpu / PreRev 0.0002165120839999 s 0.000249955752999 s 0.87
scatter_sum / HLOOpt / tpu / PostRev 0.0002144233149992 s 0.0002499135930011 s 0.86
scatter_sum / HLOOpt / tpu / BothRev 0.0002048458280005 s 0.0002362329770003 s 0.87
scatter_sum / PartOpt / tpu / PreRev 0.0002178531130011 s 0.0002356470859995 s 0.92
scatter_sum / PartOpt / tpu / PostRev 0.0002090524369996 s 0.0002204431600002 s 0.95
scatter_sum / PartOpt / tpu / BothRev 0.000225914999999 s 0.0002061643829983 s 1.10
scatter_sum / IPartOpt / tpu / PreRev 0.0002305776080029 s 0.0002229880609993 s 1.03
scatter_sum / IPartOpt / tpu / PostRev 0.0002346780270017 s 0.0002451685310006 s 0.96
scatter_sum / IPartOpt / tpu / BothRev 0.0001998067599997 s 0.0002399665489974 s 0.83
scatter_sum / DefOpt / tpu / PreRev 0.0002128472950025 s 0.0002122784159982 s 1.00
scatter_sum / DefOpt / tpu / PostRev 0.0002338053159983 s 0.0002125397060008 s 1.10
scatter_sum / DefOpt / tpu / BothRev 0.0002228634210005 s 0.0002478843019998 s 0.90
scatter_sum / IDefOpt / tpu / PreRev 0.000233180676998 s 0.0002418194800011 s 0.96
scatter_sum / IDefOpt / tpu / PostRev 0.0002271964889987 s 0.0002438022300011 s 0.93
scatter_sum / IDefOpt / tpu / BothRev 0.0002220891619981 s 0.0002428377200012 s 0.91
scatter_sum / JaXPipe / cpu / Primal 0.000007568293000076665 s 0.000005582883999977639 s 1.36
scatter_sum / Jax / cpu / Primal 0.000007597621999593684 s 0.000005537862999972276 s 1.37
scatter_sum / HLOOpt / cpu / Primal 0.000007988946001205477 s 0.000005586341999787692 s 1.43
scatter_sum / PartOpt / cpu / Primal 0.000007499318000554922 s 0.00000549940799965043 s 1.36
scatter_sum / IPartOpt / cpu / Primal 0.000007632277000084287 s 0.000005470600000080594 s 1.40
scatter_sum / DefOpt / cpu / Primal 0.000008006308000403806 s 0.000005552988999625086 s 1.44
scatter_sum / IDefOpt / cpu / Primal 0.000007671072999073659 s 0.000005579821999617707 s 1.37
scatter_sum / JaXPipe / cpu / Forward 0.000010909102000368875 s 0.000008504458000061277 s 1.28
scatter_sum / Jax / cpu / Forward 0.000011427897999965354 s 0.000008450866999737627 s 1.35
scatter_sum / HLOOpt / cpu / Forward 0.000010930751001069438 s 0.00000851859600015814 s 1.28
scatter_sum / PartOpt / cpu / Forward 0.000011513362000187043 s 0.000008492328000102133 s 1.36
scatter_sum / IPartOpt / cpu / Forward 0.000011480997998660314 s 0.000008530288999736512 s 1.35
scatter_sum / DefOpt / cpu / Forward 0.000011530764999406529 s 0.000008528337999905488 s 1.35
scatter_sum / IDefOpt / cpu / Forward 0.000011018993998732186 s 0.00000847586500003672 s 1.30
scatter_sum / JaXPipe / cpu / PreRev 0.000011617313999522592 s 0.000008587461999923108 s 1.35
scatter_sum / JaXPipe / cpu / PostRev 0.000011538565000591916 s 0.000008512100999723771 s 1.36
scatter_sum / JaXPipe / cpu / BothRev 0.000011550577999514644 s 0.000008555274000173086 s 1.35
scatter_sum / Jax / cpu / BothRev 0.000011627484000200638 s 0.000008547964000172215 s 1.36
scatter_sum / HLOOpt / cpu / PreRev 0.00001155364399892278 s 0.000008617295000021841 s 1.34
scatter_sum / HLOOpt / cpu / PostRev 0.00001098899899989192 s 0.000008582896999996592 s 1.28
scatter_sum / HLOOpt / cpu / BothRev 0.00001154438799858326 s 0.000008562118000099872 s 1.35
scatter_sum / PartOpt / cpu / PreRev 0.00001164295399939874 s 0.000008642949999739358 s 1.35
scatter_sum / PartOpt / cpu / PostRev 0.00001162131099954422 s 0.000008591775999775564 s 1.35
scatter_sum / PartOpt / cpu / BothRev 0.000011049597998862738 s 0.000008602850999977817 s 1.28
scatter_sum / IPartOpt / cpu / PreRev 0.00001156338300097559 s 0.000008546456999738439 s 1.35
scatter_sum / IPartOpt / cpu / PostRev 0.000011587959999815212 s 0.000008552622000024712 s 1.35
scatter_sum / IPartOpt / cpu / BothRev 0.000011654578000161565 s 0.000008500688999902194 s 1.37
scatter_sum / DefOpt / cpu / PreRev 0.000011645474000033572 s 0.000008667650000006688 s 1.34
scatter_sum / DefOpt / cpu / PostRev 0.000011645122000118136 s 0.000008575887000006332 s 1.36
scatter_sum / DefOpt / cpu / BothRev 0.000011590846999752102 s 0.000008634672999960457 s 1.34
scatter_sum / IDefOpt / cpu / PreRev 0.000011594143999900552 s 0.000008544630999949731 s 1.36
scatter_sum / IDefOpt / cpu / PostRev 0.000011576839000554174 s 0.000009006696000142256 s 1.29
scatter_sum / IDefOpt / cpu / BothRev 0.000011578566000025602 s 0.000008609931000137294 s 1.34
slicing / JaXPipe / cpu / Primal 0.000008656453999719815 s 0.000004426266000336909 s 1.96
slicing / Jax / cpu / Primal 0.000009221158998116152 s 0.000004467744000066887 s 2.06
slicing / HLOOpt / cpu / Primal 0.000009144225998170442 s 0.000004462859999875946 s 2.05
slicing / PartOpt / cpu / Primal 0.000009141912996710744 s 0.000004476529999919876 s 2.04
slicing / IPartOpt / cpu / Primal 0.000008684680997248506 s 0.000004420699000093009 s 1.96
slicing / DefOpt / cpu / Primal 0.000009236483001586747 s 0.000004484051999952499 s 2.06
slicing / IDefOpt / cpu / Primal 0.000008460958997602575 s 0.000004518107999956555 s 1.87
slicing / JaXPipe / cpu / Forward 0.000012548348000564149 s 0.000006957301000056759 s 1.80
slicing / Jax / cpu / Forward 0.000013235312999313464 s 0.000006877776000237645 s 1.92
slicing / HLOOpt / cpu / Forward 0.000012456705000658986 s 0.000006886585999836825 s 1.81
slicing / PartOpt / cpu / Forward 0.000013221919998613884 s 0.000006918958999904134 s 1.91
slicing / IPartOpt / cpu / Forward 0.000013298633999511368 s 0.000006928395999693748 s 1.92
slicing / DefOpt / cpu / Forward 0.00001335232700148481 s 0.000006835684000179754 s 1.95
slicing / IDefOpt / cpu / Forward 0.000013277117999678012 s 0.000006877685000290512 s 1.93
slicing / JaXPipe / cpu / PreRev 0.000014101458000368438 s 0.000007317839999814168 s 1.93
slicing / JaXPipe / cpu / PostRev 0.000013467729000694815 s 0.000007341053999880387 s 1.83
slicing / JaXPipe / cpu / BothRev 0.000014225314000213983 s 0.000007389728999896761 s 1.93
slicing / Jax / cpu / BothRev 0.000014226825998775894 s 0.000007308600000214938 s 1.95
slicing / HLOOpt / cpu / PreRev 0.000013457614000799368 s 0.00000733682600002794 s 1.83
slicing / HLOOpt / cpu / PostRev 0.000014156170000205747 s 0.000007411324000258901 s 1.91
slicing / HLOOpt / cpu / BothRev 0.000014276084999437444 s 0.000007383064000350714 s 1.93
slicing / PartOpt / cpu / PreRev 0.00001423635800165357 s 0.0000073411899998063745 s 1.94
slicing / PartOpt / cpu / PostRev 0.000013494062000972917 s 0.000007302879000235407 s 1.85
slicing / PartOpt / cpu / BothRev 0.00001344518100086134 s 0.000007329056999878958 s 1.83
slicing / IPartOpt / cpu / PreRev 0.000013481412999681195 s 0.000007322281000142538 s 1.84
slicing / IPartOpt / cpu / PostRev 0.000014188021999871126 s 0.000007356052999966778 s 1.93
slicing / IPartOpt / cpu / BothRev 0.000014244506000977708 s 0.000007383969999864348 s 1.93
slicing / DefOpt / cpu / PreRev 0.000013368766998610229 s 0.000007362848999946436 s 1.82
slicing / DefOpt / cpu / PostRev 0.000014156075998471352 s 0.000007398277999982383 s 1.91
slicing / DefOpt / cpu / BothRev 0.000014234784001018853 s 0.000007346049000261699 s 1.94
slicing / IDefOpt / cpu / PreRev 0.00001338757199846441 s 0.000007326506000026711 s 1.83
slicing / IDefOpt / cpu / PostRev 0.000014026837001438252 s 0.0000073578900000939025 s 1.91
slicing / IDefOpt / cpu / BothRev 0.000014185268999426626 s 0.000007380155000191735 s 1.92
slicing / JaXPipe / tpu / Primal 0.0001510192690002 s 0.0001480263869998 s 1.02
slicing / Jax / tpu / Primal 0.0001514681890002 s 0.0001332768899992 s 1.14
slicing / HLOOpt / tpu / Primal 0.0001483825310024 s 0.000133857621 s 1.11
slicing / PartOpt / tpu / Primal 0.0001504456500006 s 0.0001393558529998 s 1.08
slicing / IPartOpt / tpu / Primal 0.000147179831998 s 0.000141882144002 s 1.04
slicing / DefOpt / tpu / Primal 0.000147935910998 s 0.0001486796069984 s 0.99
slicing / IDefOpt / tpu / Primal 0.0001523360989995 s 0.0001484841669989 s 1.03
slicing / JaXPipe / tpu / Forward 0.0002197673619994 s 0.0002302208939981 s 0.95
slicing / Jax / tpu / Forward 0.0002331942570017 s 0.0002444545210018 s 0.95
slicing / HLOOpt / tpu / Forward 0.0002266232100009 s 0.0002368108570008 s 0.96
slicing / PartOpt / tpu / Forward 0.0002267485690026 s 0.0002269456729991 s 1.00
slicing / IPartOpt / tpu / Forward 0.0002225325609979 s 0.0002252673620023 s 0.99
slicing / DefOpt / tpu / Forward 0.0002267383800026 s 0.0002267733030021 s 1.00
slicing / IDefOpt / tpu / Forward 0.0002261034999974 s 0.0002271081619983 s 1.00
slicing / JaXPipe / tpu / PreRev 0.0002270540599965 s 0.0002302926539996 s 0.99
slicing / JaXPipe / tpu / PostRev 0.0002451873130012 s 0.000214666658001 s 1.14
slicing / JaXPipe / tpu / BothRev 0.0002368952350007 s 0.0002169244379983 s 1.09
slicing / Jax / tpu / BothRev 0.0002350982360003 s 0.0002214252199992 s 1.06
slicing / HLOOpt / tpu / PreRev 0.0002280547689988 s 0.0002247725619999 s 1.01
slicing / HLOOpt / tpu / PostRev 0.0002179976229999 s 0.0002238840120007 s 0.97
slicing / HLOOpt / tpu / BothRev 0.0002019406100007 s 0.0002118412460004 s 0.95
slicing / PartOpt / tpu / PreRev 0.0002332867570003 s 0.000202149880999 s 1.15
slicing / PartOpt / tpu / PostRev 0.0002037619689981 s 0.0002027671520008 s 1.00
slicing / PartOpt / tpu / BothRev 0.0002125789649981 s 0.000196876699003 s 1.08
slicing / IPartOpt / tpu / PreRev 0.0002384996849978 s 0.0002031497019997 s 1.17
slicing / IPartOpt / tpu / PostRev 0.000233258147 s 0.0002210606900007 s 1.06
slicing / IPartOpt / tpu / BothRev 0.0002188641520006 s 0.0002123238560016 s 1.03
slicing / DefOpt / tpu / PreRev 0.0002172725429991 s 0.0002074338339989 s 1.05
slicing / DefOpt / tpu / PostRev 0.0002297404579985 s 0.0002183732590019 s 1.05
slicing / DefOpt / tpu / BothRev 0.0002236696010004 s 0.0002048767330015 s 1.09
slicing / IDefOpt / tpu / PreRev 0.0001974652919998 s 0.0002204469400021 s 0.90
slicing / IDefOpt / tpu / PostRev 0.0002007913900015 s 0.0002183252989998 s 0.92
slicing / IDefOpt / tpu / BothRev 0.0002349731569993 s 0.0002019834810016 s 1.16
slicing / JaXPipe / cpu / Primal 0.000006140378000054625 s 0.000004426266000336909 s 1.39
slicing / Jax / cpu / Primal 0.000006518633001178386 s 0.000004467744000066887 s 1.46
slicing / HLOOpt / cpu / Primal 0.000006459558000642574 s 0.000004462859999875946 s 1.45
slicing / PartOpt / cpu / Primal 0.0000064158509994740594 s 0.000004476529999919876 s 1.43
slicing / IPartOpt / cpu / Primal 0.000006430892000935273 s 0.000004420699000093009 s 1.45
slicing / DefOpt / cpu / Primal 0.000006418057000701083 s 0.000004484051999952499 s 1.43
slicing / IDefOpt / cpu / Primal 0.000006209982000655145 s 0.000004518107999956555 s 1.37
slicing / JaXPipe / cpu / Forward 0.000008773539999310742 s 0.000006957301000056759 s 1.26
slicing / Jax / cpu / Forward 0.000008812984999167384 s 0.000006877776000237645 s 1.28
slicing / HLOOpt / cpu / Forward 0.000008815061999484896 s 0.000006886585999836825 s 1.28
slicing / PartOpt / cpu / Forward 0.000009307563999755076 s 0.000006918958999904134 s 1.35
slicing / IPartOpt / cpu / Forward 0.000009319859000243014 s 0.000006928395999693748 s 1.35
slicing / DefOpt / cpu / Forward 0.000009316614001363631 s 0.000006835684000179754 s 1.36
slicing / IDefOpt / cpu / Forward 0.000009319744000094942 s 0.000006877685000290512 s 1.36
slicing / JaXPipe / cpu / PreRev 0.000009474630998738575 s 0.000007317839999814168 s 1.29
slicing / JaXPipe / cpu / PostRev 0.000009415212000021712 s 0.000007341053999880387 s 1.28
slicing / JaXPipe / cpu / BothRev 0.000009964241000488984 s 0.000007389728999896761 s 1.35
slicing / Jax / cpu / BothRev 0.000009944892999556032 s 0.000007308600000214938 s 1.36
slicing / HLOOpt / cpu / PreRev 0.000010071315000459436 s 0.00000733682600002794 s 1.37
slicing / HLOOpt / cpu / PostRev 0.000009993196001232718 s 0.000007411324000258901 s 1.35
slicing / HLOOpt / cpu / BothRev 0.00000939696200111939 s 0.000007383064000350714 s 1.27
slicing / PartOpt / cpu / PreRev 0.000009474573998886626 s 0.0000073411899998063745 s 1.29
slicing / PartOpt / cpu / PostRev 0.000009443525999813571 s 0.000007302879000235407 s 1.29
slicing / PartOpt / cpu / BothRev 0.000009450696001295 s 0.000007329056999878958 s 1.29
slicing / IPartOpt / cpu / PreRev 0.000009401033999893116 s 0.000007322281000142538 s 1.28
slicing / IPartOpt / cpu / PostRev 0.000010024826000517351 s 0.000007356052999966778 s 1.36
slicing / IPartOpt / cpu / BothRev 0.00001002268299998832 s 0.000007383969999864348 s 1.36
slicing / DefOpt / cpu / PreRev 0.000009469299000556927 s 0.000007362848999946436 s 1.29
slicing / DefOpt / cpu / PostRev 0.000010119414999280707 s 0.000007398277999982383 s 1.37
slicing / DefOpt / cpu / BothRev 0.000010000840999055072 s 0.000007346049000261699 s 1.36
slicing / IDefOpt / cpu / PreRev 0.00000944280899966543 s 0.000007326506000026711 s 1.29
slicing / IDefOpt / cpu / PostRev 0.000009939045999999509 s 0.0000073578900000939025 s 1.35
slicing / IDefOpt / cpu / BothRev 0.00000937963199976366 s 0.000007380155000191735 s 1.27
sum / JaXPipe / cpu / Primal 0.000011808063001808478 s 0.000005717654999898514 s 2.07
sum / Jax / cpu / Primal 0.000011804295998445014 s 0.000005959596999673522 s 1.98
sum / HLOOpt / cpu / Primal 0.000012398380000377074 s 0.000005985720000353468 s 2.07
sum / PartOpt / cpu / Primal 0.00001240870299807284 s 0.000005668724999850383 s 2.19
sum / IPartOpt / cpu / Primal 0.000011897686999873256 s 0.000005971332999706646 s 1.99
sum / DefOpt / cpu / Primal 0.000012444195999705698 s 0.0000057181749998562734 s 2.18
sum / IDefOpt / cpu / Primal 0.000011674445999233285 s 0.000005982135000067501 s 1.95
sum / JaXPipe / cpu / Forward 0.000018179587001213804 s 0.00000932663200001116 s 1.95
sum / Jax / cpu / Forward 0.00001816049400076736 s 0.000009242949000054068 s 1.96
sum / HLOOpt / cpu / Forward 0.000018154380002670223 s 0.00000930015399990225 s 1.95
sum / PartOpt / cpu / Forward 0.000018221423000795768 s 0.000009327125000254456 s 1.95
sum / IPartOpt / cpu / Forward 0.00001734797599783633 s 0.000009220400999765845 s 1.88
sum / DefOpt / cpu / Forward 0.00001743864599848166 s 0.000009310176999861142 s 1.87
sum / IDefOpt / cpu / Forward 0.000017510092999145853 s 0.00000926165500004572 s 1.89
sum / JaXPipe / cpu / PreRev 0.000016903875999560112 s 0.000008530221000000892 s 1.98
sum / JaXPipe / cpu / PostRev 0.000016146939000464045 s 0.000008101518000330544 s 1.99
sum / JaXPipe / cpu / BothRev 0.000016922542999964208 s 0.000008239699000114342 s 2.05
sum / Jax / cpu / BothRev 0.00001623989499785239 s 0.000008138999000038894 s 2.00
sum / HLOOpt / cpu / PreRev 0.000017035643999406603 s 0.000008184364000044297 s 2.08
sum / HLOOpt / cpu / PostRev 0.000016921374000958166 s 0.00000816211299979841 s 2.07
sum / HLOOpt / cpu / BothRev 0.000016962646001047687 s 0.00000815452299957542 s 2.08
sum / PartOpt / cpu / PreRev 0.000016988085000775753 s 0.00000810282400016149 s 2.10
sum / PartOpt / cpu / PostRev 0.000016778936998889547 s 0.000008163558999967791 s 2.06
sum / PartOpt / cpu / BothRev 0.00001609996299885097 s 0.00000818542900015018 s 1.97
sum / IPartOpt / cpu / PreRev 0.00001603863299897057 s 0.000008087128000170197 s 1.98
sum / IPartOpt / cpu / PostRev 0.00001685486900169053 s 0.000008180826000170783 s 2.06
sum / IPartOpt / cpu / BothRev 0.00001685560199985048 s 0.000008178305999990699 s 2.06
sum / DefOpt / cpu / PreRev 0.000017093360002036205 s 0.000008128382999984752 s 2.10
sum / DefOpt / cpu / PostRev 0.000016893886000616477 s 0.000008158633999755693 s 2.07
sum / DefOpt / cpu / BothRev 0.000016035119002481224 s 0.000008214349999889237 s 1.95
sum / IDefOpt / cpu / PreRev 0.000016797018997749546 s 0.000008286450000014157 s 2.03
sum / IDefOpt / cpu / PostRev 0.000016813269001431762 s 0.000008550756000204273 s 1.97
sum / IDefOpt / cpu / BothRev 0.000016163543001312064 s 0.00000818977500011897 s 1.97
sum / JaXPipe / tpu / Primal 0.000147696791002 s 0.0001347370209987 s 1.10
sum / Jax / tpu / Primal 0.0001476192910013 s 0.0001383535230015 s 1.07
sum / HLOOpt / tpu / Primal 0.0001375304950015 s 0.0001391451830022 s 0.99
sum / PartOpt / tpu / Primal 0.000140283323999 s 0.0001342352610008 s 1.05
sum / IPartOpt / tpu / Primal 0.000144346803001 s 0.0001340883310003 s 1.08
sum / DefOpt / tpu / Primal 0.0001444204430008 s 0.0001383991319999 s 1.04
sum / IDefOpt / tpu / Primal 0.000143856421997 s 0.0001385904830021 s 1.04
sum / JaXPipe / tpu / Forward 0.0002005242099985 s 0.0001992874499992 s 1.01
sum / Jax / tpu / Forward 0.0001882580749988 s 0.0002025275619998 s 0.93
sum / HLOOpt / tpu / Forward 0.0002075211580013 s 0.0002198592699969 s 0.94
sum / PartOpt / tpu / Forward 0.0002017017300022 s 0.0002203388399975 s 0.92
sum / IPartOpt / tpu / Forward 0.0001782838289982 s 0.0002212226999981 s 0.81
sum / DefOpt / tpu / Forward 0.000204165138999 s 0.0002017273219971 s 1.01
sum / IDefOpt / tpu / Forward 0.0002241245410004 s 0.0002033278119997 s 1.10
sum / JaXPipe / tpu / PreRev 0.0001959467220003 s 0.0002247929820005 s 0.87
sum / JaXPipe / tpu / PostRev 0.0002012083100016 s 0.0002200986899988 s 0.91
sum / JaXPipe / tpu / BothRev 0.0002097789059989 s 0.0002210274299977 s 0.95
sum / Jax / tpu / BothRev 0.000202547629 s 0.0002243119920021 s 0.90
sum / HLOOpt / tpu / PreRev 0.0002204091919993 s 0.000221493959998 s 1.00
sum / HLOOpt / tpu / PostRev 0.0002196793830007 s 0.0002217695710023 s 0.99
sum / HLOOpt / tpu / BothRev 0.0002142447250007 s 0.0002253128219999 s 0.95
sum / PartOpt / tpu / PreRev 0.0002283467080014 s 0.0002210856400015 s 1.03
sum / PartOpt / tpu / PostRev 0.000226990389001 s 0.0002159380170014 s 1.05
sum / PartOpt / tpu / BothRev 0.0002091569059994 s 0.0001975620499979 s 1.06
sum / IPartOpt / tpu / PreRev 0.0002158225739985 s 0.0002151447980031 s 1.00
sum / IPartOpt / tpu / PostRev 0.0002274416489999 s 0.0002018078319997 s 1.13
sum / IPartOpt / tpu / BothRev 0.0002272989789998 s 0.000201881560999 s 1.13
sum / DefOpt / tpu / PreRev 0.0002185903230019 s 0.0002015559609972 s 1.08
sum / DefOpt / tpu / PostRev 0.0002043933790009 s 0.0002135788269988 s 0.96
sum / DefOpt / tpu / BothRev 0.0002212358819997 s 0.0002030389819992 s 1.09
sum / IDefOpt / tpu / PreRev 0.0002199346430024 s 0.0001974986390014 s 1.11
sum / IDefOpt / tpu / PostRev 0.0002038714690024 s 0.0002214387400017 s 0.92
sum / IDefOpt / tpu / BothRev 0.0001992346299994 s 0.0002198063789983 s 0.91
sum / JaXPipe / cpu / Primal 0.000008457357998850058 s 0.000005717654999898514 s 1.48
sum / Jax / cpu / Primal 0.000008143244000166305 s 0.000005959596999673522 s 1.37
sum / HLOOpt / cpu / Primal 0.000008405974000197602 s 0.000005985720000353468 s 1.40
sum / PartOpt / cpu / Primal 0.000008034278000195627 s 0.000005668724999850383 s 1.42
sum / IPartOpt / cpu / Primal 0.000008123747000354343 s 0.000005971332999706646 s 1.36
sum / DefOpt / cpu / Primal 0.000008153086000675102 s 0.0000057181749998562734 s 1.43
sum / IDefOpt / cpu / Primal 0.000008240126999226049 s 0.000005982135000067501 s 1.38
sum / JaXPipe / cpu / Forward 0.000012503977999585912 s 0.00000932663200001116 s 1.34
sum / Jax / cpu / Forward 0.000012341835001279831 s 0.000009242949000054068 s 1.34
sum / HLOOpt / cpu / Forward 0.000012462414999390604 s 0.00000930015399990225 s 1.34
sum / PartOpt / cpu / Forward 0.00001247001899901079 s 0.000009327125000254456 s 1.34
sum / IPartOpt / cpu / Forward 0.000012573812000482576 s 0.000009220400999765845 s 1.36
sum / DefOpt / cpu / Forward 0.000012441613000191863 s 0.000009310176999861142 s 1.34
sum / IDefOpt / cpu / Forward 0.000012558353000713396 s 0.00000926165500004572 s 1.36
sum / JaXPipe / cpu / PreRev 0.00001160860500021954 s 0.000008530221000000892 s 1.36
sum / JaXPipe / cpu / PostRev 0.000010970082999847363 s 0.000008101518000330544 s 1.35
sum / JaXPipe / cpu / BothRev 0.000011581274999116431 s 0.000008239699000114342 s 1.41
sum / Jax / cpu / BothRev 0.00001105305399869394 s 0.000008138999000038894 s 1.36
sum / HLOOpt / cpu / PreRev 0.000011611959998845122 s 0.000008184364000044297 s 1.42
sum / HLOOpt / cpu / PostRev 0.00001110828700075217 s 0.00000816211299979841 s 1.36
sum / HLOOpt / cpu / BothRev 0.000011602652999499696 s 0.00000815452299957542 s 1.42
sum / PartOpt / cpu / PreRev 0.000011047512998629829 s 0.00000810282400016149 s 1.36
sum / PartOpt / cpu / PostRev 0.0000116581110014522 s 0.000008163558999967791 s 1.43
sum / PartOpt / cpu / BothRev 0.000011113969001598887 s 0.00000818542900015018 s 1.36
sum / IPartOpt / cpu / PreRev 0.000011557420999452008 s 0.000008087128000170197 s 1.43
sum / IPartOpt / cpu / PostRev 0.00001102848200025619 s 0.000008180826000170783 s 1.35
sum / IPartOpt / cpu / BothRev 0.000011688472999594525 s 0.000008178305999990699 s 1.43
sum / DefOpt / cpu / PreRev 0.00001156935600010911 s 0.000008128382999984752 s 1.42
sum / DefOpt / cpu / PostRev 0.000011604893999901834 s 0.000008158633999755693 s 1.42
sum / DefOpt / cpu / BothRev 0.000011466003999885287 s 0.000008214349999889237 s 1.40
sum / IDefOpt / cpu / PreRev 0.00001166120299967588 s 0.000008286450000014157 s 1.41
sum / IDefOpt / cpu / PostRev 0.000011637196999799923 s 0.000008550756000204273 s 1.36
sum / IDefOpt / cpu / BothRev 0.000011662502000035602 s 0.00000818977500011897 s 1.42
value_and_grad / JaXPipe / cpu / Primal 0.000017053715000656667 s 0.000009939150999798585 s 1.72
value_and_grad / Jax / cpu / Primal 0.000017932528000528692 s 0.000010449653000250692 s 1.72
value_and_grad / HLOOpt / cpu / Primal 0.000018035844001133223 s 0.000010490282999853662 s 1.72
value_and_grad / PartOpt / cpu / Primal 0.000018339242997171825 s 0.00001061113000014302 s 1.73
value_and_grad / IPartOpt / cpu / Primal 0.000018150171999877785 s 0.000010577111999737098 s 1.72
value_and_grad / DefOpt / cpu / Primal 0.00001820260499880533 s 0.000010521465000238094 s 1.73
value_and_grad / IDefOpt / cpu / Primal 0.00001842676399974153 s 0.00001074638199997935 s 1.71
value_and_grad / JaXPipe / tpu / Primal 0.0002298873679974 s 0.000231804594001 s 0.99
value_and_grad / Jax / tpu / Primal 0.0002381845449999 s 0.0002095536649976 s 1.14
value_and_grad / HLOOpt / tpu / Primal 0.000236947835001 s 0.0002412452300013 s 0.98
value_and_grad / PartOpt / tpu / Primal 0.000237230775001 s 0.0002437919900003 s 0.97
value_and_grad / IPartOpt / tpu / Primal 0.0002333642070007 s 0.0002172322280021 s 1.07
value_and_grad / DefOpt / tpu / Primal 0.0002390224840019 s 0.0002067109130002 s 1.16
value_and_grad / IDefOpt / tpu / Primal 0.0002514215590017 s 0.0002449896010002 s 1.03
value_and_grad / JaXPipe / cpu / Primal 0.00001238375899993116 s 0.000009939150999798585 s 1.25
value_and_grad / Jax / cpu / Primal 0.000012264216000403392 s 0.000010449653000250692 s 1.17
value_and_grad / HLOOpt / cpu / Primal 0.000012961440999788463 s 0.000010490282999853662 s 1.24
value_and_grad / PartOpt / cpu / Primal 0.000012941006998516968 s 0.00001061113000014302 s 1.22
value_and_grad / IPartOpt / cpu / Primal 0.000012697979000222404 s 0.000010577111999737098 s 1.20
value_and_grad / DefOpt / cpu / Primal 0.00001283516600051371 s 0.000010521465000238094 s 1.22
value_and_grad / IDefOpt / cpu / Primal 0.00001293278699995426 s 0.00001074638199997935 s 1.20
llama / JaXPipe / tpu / Primal 0.0003630122599861 s 0.0003664719800144 s 0.99
llama / Jax / tpu / Primal 0.0003616372599935 s 0.0003762115800054 s 0.96
llama / HLOOpt / tpu / Primal 0.0003604270600044 s 0.0003746317799959 s 0.96
llama / PartOpt / tpu / Primal 0.0003726020600151 s 0.0003897355600201 s 0.96
llama / IPartOpt / tpu / Primal 0.0003621772600308 s 0.0003917921799438 s 0.92
llama / DefOpt / tpu / Primal 0.0003271898799721 s 0.0003499631600425 s 0.93
llama / IDefOpt / tpu / Primal 0.0003410836799594 s 0.0003767341600178 s 0.91
llama / JaXPipe / tpu / Forward 0.0005514091800432 s 0.0005924258600134 s 0.93
llama / Jax / tpu / Forward 0.0006798613400314 s 0.0007164409200049 s 0.95
llama / HLOOpt / tpu / Forward 0.000546684180008 s 0.0005895116600004 s 0.93
llama / PartOpt / tpu / Forward 0.0005397690000245 s 0.0005847404799715 s 0.92
llama / IPartOpt / tpu / Forward 0.0005477911799971 s 0.000583398259987 s 0.94
llama / DefOpt / tpu / Forward 0.0005528787799994 s 0.0005617104600241 s 0.98
llama / IDefOpt / tpu / Forward 0.0005610035800054 s 0.000567741440027 s 0.99
llama / JaXPipe / tpu / PreRev 0.0007950609000545 s 0.000761423340009 s 1.04
llama / JaXPipe / tpu / PostRev 0.0007460639000055 s 0.0007672383599856 s 0.97
llama / JaXPipe / tpu / BothRev 0.0007833812999888 s 0.000804779980026 s 0.97
llama / Jax / tpu / BothRev 0.0007514615199761 s 0.0007698797400371 s 0.98
llama / HLOOpt / tpu / PreRev 0.0007980478999525 s 0.0008063667600072 s 0.99
llama / HLOOpt / tpu / PostRev 0.0007856474800064 s 0.0008024775599915 s 0.98
llama / HLOOpt / tpu / BothRev 0.0007835178999812 s 0.0007608395400166 s 1.03
llama / PartOpt / tpu / PreRev 0.000787046080004 s 0.0007604935400013 s 1.03
llama / PartOpt / tpu / PostRev 0.0007634269000118 s 0.0007338385400362 s 1.04
llama / PartOpt / tpu / BothRev 0.000787314620029 s 0.000751907340018 s 1.05
llama / IPartOpt / tpu / PreRev 0.0007864662400243 s 0.0007484921400464 s 1.05
llama / IPartOpt / tpu / PostRev 0.0007405916399875 s 0.0007285823199345 s 1.02
llama / IPartOpt / tpu / BothRev 0.0007960162399831 s 0.0007437195400416 s 1.07
llama / DefOpt / tpu / PreRev 0.0007892280399391 s 0.0007607689399446 s 1.04
llama / DefOpt / tpu / PostRev 0.0007425096399674 s 0.0007955995600059 s 0.93
llama / DefOpt / tpu / BothRev 0.0007855168799869 s 0.0008063657800084 s 0.97
llama / IDefOpt / tpu / PreRev 0.0007877307000308 s 0.0008002043799933 s 0.98
llama / IDefOpt / tpu / PostRev 0.0007505319199844 s 0.0007632747600291 s 0.98
llama / IDefOpt / tpu / BothRev 0.0007840082999609 s 0.00076530696002 s 1.02
llama / JaXPipe / cpu / Primal 0.0014763623999897 s 0.0008527092999884 s 1.73
llama / Jax / cpu / Primal 0.0014259253001 s 0.0008616218000042 s 1.65
llama / HLOOpt / cpu / Primal 0.0015814403001058 s 0.0009334207999927 s 1.69
llama / PartOpt / cpu / Primal 0.0015311331000702 s 0.000859156899969 s 1.78
llama / IPartOpt / cpu / Primal 0.0014123605000349 s 0.0008582037999985 s 1.65
llama / DefOpt / cpu / Primal 0.0015591119999953 s 0.0009581214000263 s 1.63
llama / IDefOpt / cpu / Primal 0.0015330592999816 s 0.0009128119000251 s 1.68
llama / JaXPipe / cpu / Forward 0.0052645658999608 s 0.0023119594000036 s 2.28
llama / Jax / cpu / Forward 0.0053721603000667 s 0.002347239599976 s 2.29
llama / HLOOpt / cpu / Forward 0.0052139445999273 s 0.0023568167999655 s 2.21
llama / PartOpt / cpu / Forward 0.0051344381001399 s 0.0024163317999864 s 2.12
llama / IPartOpt / cpu / Forward 0.0050258032999408 s 0.0023457631999917 s 2.14
llama / DefOpt / cpu / Forward 0.0051159757000277 s 0.0023569691999909 s 2.17
llama / IDefOpt / cpu / Forward 0.0051930567999079 s 0.0023619774000053 s 2.20
llama / JaXPipe / cpu / PreRev 0.0095432697000433 s 0.0061528604000159 s 1.55
llama / JaXPipe / cpu / PostRev 0.0098230739999053 s 0.0054835653999816 s 1.79
llama / JaXPipe / cpu / BothRev 0.0095235360000515 s 0.0049855954000122 s 1.91
llama / Jax / cpu / BothRev 0.0087260405998677 s 0.0055081307999898 s 1.58
llama / HLOOpt / cpu / PreRev 0.0099959160999787 s 0.0049616335999871 s 2.01
llama / HLOOpt / cpu / PostRev 0.0066579103000549 s 0.005198626099991 s 1.28
llama / HLOOpt / cpu / BothRev 0.0097886798001127 s 0.0055855210000117 s 1.75
llama / PartOpt / cpu / PreRev 0.0075892201999522 s 0.0048668442000234 s 1.56
llama / PartOpt / cpu / PostRev 0.0097966227998767 s 0.0055047661000116 s 1.78
llama / PartOpt / cpu / BothRev 0.0094534038000347 s 0.0048355208999964 s 1.95
llama / IPartOpt / cpu / PreRev 0.0104731655999785 s 0.0054107675999603 s 1.94
llama / IPartOpt / cpu / PostRev 0.0070844493999175 s 0.0048973619000207 s 1.45
llama / IPartOpt / cpu / BothRev 0.0102373715999419 s 0.0053043362999687 s 1.93
llama / DefOpt / cpu / PreRev 0.0092598035000264 s 0.0055278509000345 s 1.68
llama / DefOpt / cpu / PostRev 0.0099568663999889 s 0.0047026706999986 s 2.12
llama / DefOpt / cpu / BothRev 0.0095541384000171 s 0.0055240848000266 s 1.73
llama / IDefOpt / cpu / PreRev 0.0104807808998884 s 0.0049009266000211 s 2.14
llama / IDefOpt / cpu / PostRev 0.0077219713999511 s 0.0054996454000047 s 1.40
llama / IDefOpt / cpu / BothRev 0.0102998726000805 s 0.0049239127999953 s 2.09

This comment was automatically generated by workflow using github-action-benchmark.

@ivanradanov ivanradanov mentioned this pull request Dec 17, 2025
Closed
wsmoses
wsmoses reviewed Dec 19, 2025
View reviewed changes
src/enzyme_ad/jax/Passes/RemoveAtomics.cpp
@@ -0,0 +1,461 @@
//===---------------------------------------------------------------------===//
Copy link
Member

@wsmoses wsmoses Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we get rid of the pollydebug.cpp file?

@ivanradanov @wsmoses
wip remove atomics
56b5ed5 
wip

wip

testing

wip

wip

done?

file

remove unneeded files

actually run test

revert isl

dep dialects

missingheader
@ivanradanov @wsmoses
stray file
c1992d4
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments

@wsmoses wsmoses wsmoses approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ivanradanov @wsmoses