Skip to content

Use mlir utils from Enzyme #1785

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
Closed
ivanradanov wants to merge 3 commits into from
Closed

Use mlir utils from Enzyme #1785

ivanradanov wants to merge 3 commits into from

Conversation

@ivanradanov
Copy link
Collaborator

@ivanradanov ivanradanov commented Dec 17, 2025

No description provided.

@ivanradanov ivanradanov requested a review from wsmoses December 17, 2025 10:28
@ivanradanov
Copy link
Collaborator Author

ivanradanov commented Dec 17, 2025

@wsmoses These are the changes that may be interfering with the tests from here #1782 but they seem to be passing?

github-actions[bot]
github-actions bot reviewed Dec 17, 2025
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment

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: ce2036d Previous: f498f06 Ratio
actmtch / JaX / cpu / Primal 0.000003898383500018099 s 0.000004198181699985071 s 0.93
actmtch / JaXPipe / cpu / Primal 0.000003889334999985295 s 0.000003863652500012904 s 1.01
actmtch / JaX / cpu / Forward 0.000006071556599999895 s 0.00000637123760002396 s 0.95
actmtch / JaXPipe / cpu / Forward 0.000007245763599985366 s 0.000007263626399981149 s 1.00
actmtch / JaX / cpu / BothRev 0.000005988591300001644 s 0.0000063148874000034994 s 0.95
actmtch / JaXPipe / cpu / PreRev 0.000007072921799999676 s 0.000007087604700018346 s 1.00
actmtch / JaXPipe / cpu / PostRev 0.0000062575698999808085 s 0.000006282340799998565 s 1.00
actmtch / JaXPipe / cpu / BothRev 0.000007090167099977407 s 0.000007413655400023345 s 0.96
actmtch / JaX / cpu / Primal 0.000007183445899863727 s 0.000004198181699985071 s 1.71
actmtch / JaXPipe / cpu / Primal 0.000007214633500552736 s 0.000003863652500012904 s 1.87
actmtch / JaX / cpu / Forward 0.000010304637299850582 s 0.00000637123760002396 s 1.62
actmtch / JaXPipe / cpu / Forward 0.000017090355698019267 s 0.000007263626399981149 s 2.35
actmtch / JaX / cpu / BothRev 0.000011095017602201551 s 0.0000063148874000034994 s 1.76
actmtch / JaXPipe / cpu / PreRev 0.000013404604099923745 s 0.000007087604700018346 s 1.89
actmtch / JaXPipe / cpu / PostRev 0.000010829867300344631 s 0.000006282340799998565 s 1.72
actmtch / JaXPipe / cpu / BothRev 0.00001533411749987863 s 0.000007413655400023345 s 2.07
actmtch / JaX / gpu / Primal 0.00007271732410008553 s 0.00007541810370003077 s 0.96
actmtch / JaXPipe / gpu / Primal 0.00007363255399977789 s 0.00007496135269993828 s 0.98
actmtch / JaX / gpu / Forward 0.0001030304033018 s 0.0001045130106998 s 0.99
actmtch / JaXPipe / gpu / Forward 0.0001041215540986 s 0.0001070459729 s 0.97
actmtch / JaX / gpu / BothRev 0.0001121186394011 s 0.0001050153675001 s 1.07
actmtch / JaXPipe / gpu / PreRev 0.0001029053870006 s 0.0001028680055998 s 1.00
actmtch / JaXPipe / gpu / PostRev 0.0001012219277006 s 0.000103167493 s 0.98
actmtch / JaXPipe / gpu / BothRev 0.000114353601 s 0.0001141021014 s 1.00
actmtch / JaX / cpu / Primal 0.000003971037500014063 s 0.000004198181699985071 s 0.95
actmtch / JaXPipe / cpu / Primal 0.000003643013600230916 s 0.000003863652500012904 s 0.94
actmtch / JaX / cpu / Forward 0.000005409416900147335 s 0.00000637123760002396 s 0.85
actmtch / JaXPipe / cpu / Forward 0.000006596272499882616 s 0.000007263626399981149 s 0.91
actmtch / JaX / cpu / BothRev 0.000005587017799916794 s 0.0000063148874000034994 s 0.88
actmtch / JaXPipe / cpu / PreRev 0.000006297788599840715 s 0.000007087604700018346 s 0.89
actmtch / JaXPipe / cpu / PostRev 0.000005632318800053327 s 0.000006282340799998565 s 0.90
actmtch / JaXPipe / cpu / BothRev 0.000006288073599716882 s 0.000007413655400023345 s 0.85
actmtch / JaX / tpu / Primal 0.0001561620961998 s 0.0001522344173004 s 1.03
actmtch / JaXPipe / tpu / Primal 0.0001478679575 s 0.0001460869764006 s 1.01
actmtch / JaX / tpu / Forward 0.0002047694462999 s 0.0002339157574999 s 0.88
actmtch / JaXPipe / tpu / Forward 0.0002076048241997 s 0.0002329354175002 s 0.89
actmtch / JaX / tpu / BothRev 0.0002111536798001 s 0.0002287737605001 s 0.92
actmtch / JaXPipe / tpu / PreRev 0.0002119963835 s 0.0002138487846998 s 0.99
actmtch / JaXPipe / tpu / PostRev 0.0002196134564997 s 0.0002233515385989 s 0.98
actmtch / JaXPipe / tpu / BothRev 0.0002226784752998 s 0.0002211717425991 s 1.01
actmtch / JaX / cpu / Primal 0.000004894325400073285 s 0.000004198181699985071 s 1.17
actmtch / JaXPipe / cpu / Primal 0.00000504814810001335 s 0.000003863652500012904 s 1.31
actmtch / JaX / cpu / Forward 0.000007429940700058069 s 0.00000637123760002396 s 1.17
actmtch / JaXPipe / cpu / Forward 0.000009153034199971445 s 0.000007263626399981149 s 1.26
actmtch / JaX / cpu / BothRev 0.000007768261399996845 s 0.0000063148874000034994 s 1.23
actmtch / JaXPipe / cpu / PreRev 0.000009295212600045489 s 0.000007087604700018346 s 1.31
actmtch / JaXPipe / cpu / PostRev 0.000007948414600014076 s 0.000006282340799998565 s 1.27
actmtch / JaXPipe / cpu / BothRev 0.000009258783699988272 s 0.000007413655400023345 s 1.25
add_one / JaX / cpu / Primal 0.00000410932770000727 s 0.000004198862800012648 s 0.98
add_one / JaXPipe / cpu / Primal 0.0000041106918999958 s 0.000004270037700007378 s 0.96
add_one / JaX / cpu / Forward 0.000006910398199988776 s 0.000007349169500002972 s 0.94
add_one / JaXPipe / cpu / Forward 0.000006920695700000578 s 0.000007493103599972528 s 0.92
add_one / JaX / cpu / BothRev 0.000007345947800013164 s 0.000007407588200021564 s 0.99
add_one / JaXPipe / cpu / PreRev 0.000007015841499969611 s 0.000007457824600032836 s 0.94
add_one / JaXPipe / cpu / PostRev 0.00000698159220000889 s 0.000007667493300004935 s 0.91
add_one / JaXPipe / cpu / BothRev 0.0000069288477000100105 s 0.000007396029700021245 s 0.94
add_one / JaX / cpu / Primal 0.00000745164540130645 s 0.000004198862800012648 s 1.77
add_one / JaXPipe / cpu / Primal 0.00000731776729808189 s 0.000004270037700007378 s 1.71
add_one / JaX / cpu / Forward 0.000011512868499266917 s 0.000007349169500002972 s 1.57
add_one / JaXPipe / cpu / Forward 0.000012115515099139884 s 0.000007493103599972528 s 1.62
add_one / JaX / cpu / BothRev 0.000011589680501492694 s 0.000007407588200021564 s 1.56
add_one / JaXPipe / cpu / PreRev 0.000012325848001637495 s 0.000007457824600032836 s 1.65
add_one / JaXPipe / cpu / PostRev 0.000012471212001401 s 0.000007667493300004935 s 1.63
add_one / JaXPipe / cpu / BothRev 0.00001163346859975718 s 0.000007396029700021245 s 1.57
add_one / JaX / gpu / Primal 0.0000758985943015432 s 0.00007734811610007455 s 0.98
add_one / JaXPipe / gpu / Primal 0.00007554895840003155 s 0.00007582890539997607 s 1.00
add_one / JaX / gpu / Forward 0.0001036315534991 s 0.0001041033552999 s 1.00
add_one / JaXPipe / gpu / Forward 0.0001040464434976 s 0.0001045790947 s 0.99
add_one / JaX / gpu / BothRev 0.0001070596011995 s 0.0001344240716 s 0.80
add_one / JaXPipe / gpu / PreRev 0.0001222817973 s 0.000114758685 s 1.07
add_one / JaXPipe / gpu / PostRev 0.0001069141993997 s 0.0001082768000998 s 0.99
add_one / JaXPipe / gpu / BothRev 0.0001207659393985 s 0.0001086698432 s 1.11
add_one / JaX / cpu / Primal 0.000003894867499911925 s 0.000004198862800012648 s 0.93
add_one / JaXPipe / cpu / Primal 0.000003857860499920207 s 0.000004270037700007378 s 0.90
add_one / JaX / cpu / Forward 0.000005884823700034758 s 0.000007349169500002972 s 0.80
add_one / JaXPipe / cpu / Forward 0.00000616741660014668 s 0.000007493103599972528 s 0.82
add_one / JaX / cpu / BothRev 0.000006243938500119839 s 0.000007407588200021564 s 0.84
add_one / JaXPipe / cpu / PreRev 0.000006263587600187748 s 0.000007457824600032836 s 0.84
add_one / JaXPipe / cpu / PostRev 0.000006243666599766584 s 0.000007667493300004935 s 0.81
add_one / JaXPipe / cpu / BothRev 0.000006256048599971109 s 0.000007396029700021245 s 0.85
add_one / JaX / tpu / Primal 0.0001409099851 s 0.0001673230471991 s 0.84
add_one / JaXPipe / tpu / Primal 0.0001374586724999 s 0.0001471067744991 s 0.93
add_one / JaX / tpu / Forward 0.0002165321387001 s 0.0002279286256001 s 0.95
add_one / JaXPipe / tpu / Forward 0.0002195422685999 s 0.0002235035506004 s 0.98
add_one / JaX / tpu / BothRev 0.0002296230495998 s 0.0002405719884001 s 0.95
add_one / JaXPipe / tpu / PreRev 0.0002311663370001 s 0.000227359636499 s 1.02
add_one / JaXPipe / tpu / PostRev 0.0002226724313 s 0.0002271916045996 s 0.98
add_one / JaXPipe / tpu / BothRev 0.0002294212147 s 0.0002311396404998 s 0.99
add_one / JaX / cpu / Primal 0.000005335182499948133 s 0.000004198862800012648 s 1.27
add_one / JaXPipe / cpu / Primal 0.000005160618699937913 s 0.000004270037700007378 s 1.21
add_one / JaX / cpu / Forward 0.000008619849299975613 s 0.000007349169500002972 s 1.17
add_one / JaXPipe / cpu / Forward 0.00000879986519994418 s 0.000007493103599972528 s 1.17
add_one / JaX / cpu / BothRev 0.000008427721899988683 s 0.000007407588200021564 s 1.14
add_one / JaXPipe / cpu / PreRev 0.00000893650650004929 s 0.000007457824600032836 s 1.20
add_one / JaXPipe / cpu / PostRev 0.000008915626700036228 s 0.000007667493300004935 s 1.16
add_one / JaXPipe / cpu / BothRev 0.000008918952499971056 s 0.000007396029700021245 s 1.21
add_two / JaX / cpu / Primal 0.000004368313899976783 s 0.000004385326700003134 s 1.00
add_two / JaXPipe / cpu / Primal 0.000004361218599979111 s 0.000004415645900007803 s 0.99
add_two / JaX / cpu / Forward 0.000007030999900007373 s 0.00000729380709999532 s 0.96
add_two / JaXPipe / cpu / Forward 0.000007081347799976356 s 0.000007372629800011055 s 0.96
add_two / JaX / cpu / BothRev 0.000008502698100028283 s 0.00000948501870002474 s 0.90
add_two / JaXPipe / cpu / PreRev 0.00000856361520000064 s 0.000009147914400000445 s 0.94
add_two / JaXPipe / cpu / PostRev 0.000008665260499992655 s 0.00000964294410000548 s 0.90
add_two / JaXPipe / cpu / BothRev 0.000008746064399974784 s 0.000009278358000028676 s 0.94
add_two / JaX / cpu / Primal 0.00000959144129883498 s 0.000004385326700003134 s 2.19
add_two / JaXPipe / cpu / Primal 0.000009079264200408943 s 0.000004415645900007803 s 2.06
add_two / JaX / cpu / Forward 0.00001184612079814542 s 0.00000729380709999532 s 1.62
add_two / JaXPipe / cpu / Forward 0.000011940821100142784 s 0.000007372629800011055 s 1.62
add_two / JaX / cpu / BothRev 0.000013802027501515113 s 0.00000948501870002474 s 1.46
add_two / JaXPipe / cpu / PreRev 0.000014659398299409076 s 0.000009147914400000445 s 1.60
add_two / JaXPipe / cpu / PostRev 0.000016210070700617506 s 0.00000964294410000548 s 1.68
add_two / JaXPipe / cpu / BothRev 0.000014558675198350102 s 0.000009278358000028676 s 1.57
add_two / JaX / gpu / Primal 0.0000856637793011032 s 0.00008004563680005958 s 1.07
add_two / JaXPipe / gpu / Primal 0.00008671708299953025 s 0.00007762399020011798 s 1.12
add_two / JaX / gpu / Forward 0.0001029280924005 s 0.0001045923401999 s 0.98
add_two / JaXPipe / gpu / Forward 0.000107190287701 s 0.0001089025867999 s 0.98
add_two / JaX / gpu / BothRev 0.0001243360024003 s 0.0001216681032999 s 1.02
add_two / JaXPipe / gpu / PreRev 0.0001205500989977 s 0.0001217656914999 s 0.99
add_two / JaXPipe / gpu / PostRev 0.0001198427527997 s 0.0001225567402998 s 0.98
add_two / JaXPipe / gpu / BothRev 0.0001189782537985 s 0.0001217016818998 s 0.98
add_two / JaX / cpu / Primal 0.0000039202534000651215 s 0.000004385326700003134 s 0.89
add_two / JaXPipe / cpu / Primal 0.000003936728500048048 s 0.000004415645900007803 s 0.89
add_two / JaX / cpu / Forward 0.000006082939699990675 s 0.00000729380709999532 s 0.83
add_two / JaXPipe / cpu / Forward 0.000006094199599829153 s 0.000007372629800011055 s 0.83
add_two / JaX / cpu / BothRev 0.000007333218099665828 s 0.00000948501870002474 s 0.77
add_two / JaXPipe / cpu / PreRev 0.000007009818200094742 s 0.000009147914400000445 s 0.77
add_two / JaXPipe / cpu / PostRev 0.000007362480099982349 s 0.00000964294410000548 s 0.76
add_two / JaXPipe / cpu / BothRev 0.000007322496100096032 s 0.000009278358000028676 s 0.79
add_two / JaX / tpu / Primal 0.0001491628759999 s 0.0001393967755997 s 1.07
add_two / JaXPipe / tpu / Primal 0.0001554906424 s 0.0001474509203995 s 1.05
add_two / JaX / tpu / Forward 0.0002259528351001 s 0.0002287422076013 s 0.99
add_two / JaXPipe / tpu / Forward 0.0002334324760999 s 0.0002290770945997 s 1.02
add_two / JaX / tpu / BothRev 0.0002479808664 s 0.0002523867583004 s 0.98
add_two / JaXPipe / tpu / PreRev 0.0002387069661002 s 0.0002401663525 s 0.99
add_two / JaXPipe / tpu / PostRev 0.0002284717831 s 0.0002403527524002 s 0.95
add_two / JaXPipe / tpu / BothRev 0.0002313957459002 s 0.0002433409174002 s 0.95
add_two / JaX / cpu / Primal 0.000005534035199980281 s 0.000004385326700003134 s 1.26
add_two / JaXPipe / cpu / Primal 0.000005409660900022573 s 0.000004415645900007803 s 1.23
add_two / JaX / cpu / Forward 0.000008638527399943997 s 0.00000729380709999532 s 1.18
add_two / JaXPipe / cpu / Forward 0.000009239194200017664 s 0.000007372629800011055 s 1.25
add_two / JaX / cpu / BothRev 0.000010553575999983875 s 0.00000948501870002474 s 1.11
add_two / JaXPipe / cpu / PreRev 0.00001052636189997429 s 0.000009147914400000445 s 1.15
add_two / JaXPipe / cpu / PostRev 0.0000105161825000323 s 0.00000964294410000548 s 1.09
add_two / JaXPipe / cpu / BothRev 0.000010267184099939186 s 0.000009278358000028676 s 1.11
cache / JaX / cpu / Primal 0.000003973087399981523 s 0.000003864575199986576 s 1.03
cache / JaXPipe / cpu / Primal 0.000004203880099976231 s 0.000003858660700007022 s 1.09
cache / JaX / cpu / Forward 0.000010203806199979226 s 0.00001055746110000655 s 0.97
cache / JaXPipe / cpu / Forward 0.00001020545029996356 s 0.00001049885870002072 s 0.97
cache / JaX / cpu / BothRev 0.000014462816199966256 s 0.000015660578500001066 s 0.92
cache / JaXPipe / cpu / PreRev 0.00001104274119998081 s 0.000011378534699997544 s 0.97
cache / JaXPipe / cpu / PostRev 0.00001470840249999128 s 0.000015280421099987507 s 0.96
cache / JaXPipe / cpu / BothRev 0.00001077101880000555 s 0.000011242050499959078 s 0.96
cache / JaX / cpu / Primal 0.000007468902098480612 s 0.000003864575199986576 s 1.93
cache / JaXPipe / cpu / Primal 0.000007440122502157465 s 0.000003858660700007022 s 1.93
cache / JaX / cpu / Forward 0.000012948128901189194 s 0.00001055746110000655 s 1.23
cache / JaXPipe / cpu / Forward 0.00001298054869985208 s 0.00001049885870002072 s 1.24
cache / JaX / cpu / BothRev 0.00002261154949956108 s 0.000015660578500001066 s 1.44
cache / JaXPipe / cpu / PreRev 0.000016868823498953133 s 0.000011378534699997544 s 1.48
cache / JaXPipe / cpu / PostRev 0.000022212952602421864 s 0.000015280421099987507 s 1.45
cache / JaXPipe / cpu / BothRev 0.00001855379900080152 s 0.000011242050499959078 s 1.65
cache / JaX / gpu / Primal 0.00007333266150089912 s 0.00007409637430009753 s 0.99
cache / JaXPipe / gpu / Primal 0.0000773553353996249 s 0.00007343267770011152 s 1.05
cache / JaX / gpu / Forward 0.00009808222169813236 s 0.0001002098614999 s 0.98
cache / JaXPipe / gpu / Forward 0.00009864515890076292 s 0.0001099262027 s 0.90
cache / JaX / gpu / BothRev 0.0001026109077996 s 0.0001064908809999 s 0.96
cache / JaXPipe / gpu / PreRev 0.0001105622194998 s 0.0001094545573001 s 1.01
cache / JaXPipe / gpu / PostRev 0.000112202083398 s 0.0001066924905 s 1.05
cache / JaXPipe / gpu / BothRev 0.0001131286867021 s 0.0001069219834 s 1.06
cache / JaX / cpu / Primal 0.000003805926499990164 s 0.000003864575199986576 s 0.98
cache / JaXPipe / cpu / Primal 0.000003910366499985685 s 0.000003858660700007022 s 1.01
cache / JaX / cpu / Forward 0.000025388150199796656 s 0.00001055746110000655 s 2.40
cache / JaXPipe / cpu / Forward 0.000026133126800050378 s 0.00001049885870002072 s 2.49
cache / JaX / cpu / BothRev 0.00003272757719969377 s 0.000015660578500001066 s 2.09
cache / JaXPipe / cpu / PreRev 0.000030857118999847444 s 0.000011378534699997544 s 2.71
cache / JaXPipe / cpu / PostRev 0.000032692405299894745 s 0.000015280421099987507 s 2.14
cache / JaXPipe / cpu / BothRev 0.000028849935800099047 s 0.000011242050499959078 s 2.57
cache / JaX / tpu / Primal 0.0001407690341999 s 0.0001508992653005 s 0.93
cache / JaXPipe / tpu / Primal 0.0001405092373002 s 0.0001527997062992 s 0.92
cache / JaX / tpu / Forward 0.0002100335341998 s 0.0002281159746009 s 0.92
cache / JaXPipe / tpu / Forward 0.0002293465596998 s 0.0002243557596011 s 1.02
cache / JaX / tpu / BothRev 0.0002182898009999 s 0.0002140056157004 s 1.02
cache / JaXPipe / tpu / PreRev 0.0002139042827002 s 0.0002191726656004 s 0.98
cache / JaXPipe / tpu / PostRev 0.0002287231328999 s 0.0002270020376003 s 1.01
cache / JaXPipe / tpu / BothRev 0.0002347950335002 s 0.0002215035045999 s 1.06
cache / JaX / cpu / Primal 0.000004973023899947293 s 0.000003864575199986576 s 1.29
cache / JaXPipe / cpu / Primal 0.000005010526100068091 s 0.000003858660700007022 s 1.30
cache / JaX / cpu / Forward 0.00002396317430002455 s 0.00001055746110000655 s 2.27
cache / JaXPipe / cpu / Forward 0.0000245992622999438 s 0.00001049885870002072 s 2.34
cache / JaX / cpu / BothRev 0.0000322837423999772 s 0.000015660578500001066 s 2.06
cache / JaXPipe / cpu / PreRev 0.000028615346100013997 s 0.000011378534699997544 s 2.51
cache / JaXPipe / cpu / PostRev 0.00003349777749999703 s 0.000015280421099987507 s 2.19
cache / JaXPipe / cpu / BothRev 0.00002613198340004601 s 0.000011242050499959078 s 2.32
Concat / JaX / cpu / Primal 0.000004090631700000813 s 0.00000410532279997824 s 1.00
Concat / JaXPipe / cpu / Primal 0.000004190779099963038 s 0.000004153241899985005 s 1.01
Concat / JaX / cpu / Forward 0.00000666832599999907 s 0.000007026319200031139 s 0.95
Concat / JaXPipe / cpu / Forward 0.000006696450500021456 s 0.000007335820499974943 s 0.91
Concat / JaX / cpu / BothRev 0.000007077351900034045 s 0.00000748516550002023 s 0.95
Concat / JaXPipe / cpu / PreRev 0.000006939439999996466 s 0.000007412402799991469 s 0.94
Concat / JaXPipe / cpu / PostRev 0.000007214731500016569 s 0.000007310956699984672 s 0.99
Concat / JaXPipe / cpu / BothRev 0.000007067817999995896 s 0.000007627980400002343 s 0.93
Concat / JaX / cpu / Primal 0.0000074507192010059955 s 0.00000410532279997824 s 1.81
Concat / JaXPipe / cpu / Primal 0.00000776477290201001 s 0.000004153241899985005 s 1.87
Concat / JaX / cpu / Forward 0.00001227639369899407 s 0.000007026319200031139 s 1.75
Concat / JaXPipe / cpu / Forward 0.000012405246499110944 s 0.000007335820499974943 s 1.69
Concat / JaX / cpu / BothRev 0.000012233466599718669 s 0.00000748516550002023 s 1.63
Concat / JaXPipe / cpu / PreRev 0.0000144450855004834 s 0.000007412402799991469 s 1.95
Concat / JaXPipe / cpu / PostRev 0.000013122607101104222 s 0.000007310956699984672 s 1.79
Concat / JaXPipe / cpu / BothRev 0.000012913027699687518 s 0.000007627980400002343 s 1.69
Concat / JaX / gpu / Primal 0.0000811357907979982 s 0.00007730801140005497 s 1.05
Concat / JaXPipe / gpu / Primal 0.00008901359500014224 s 0.00007747235530005128 s 1.15
Concat / JaX / gpu / Forward 0.0001172159912995 s 0.0001102012127001 s 1.06
Concat / JaXPipe / gpu / Forward 0.0001068036011012 s 0.00010418491 s 1.03
Concat / JaX / gpu / BothRev 0.0001034000604005 s 0.00009770456010010094 s 1.06
Concat / JaXPipe / gpu / PreRev 0.0000975791619013762 s 0.00009882451080011378 s 0.99
Concat / JaXPipe / gpu / PostRev 0.00009730000539857428 s 0.00009880285050003294 s 0.98
Concat / JaXPipe / gpu / BothRev 0.0001214398425014 s 0.00009874160970011872 s 1.23
Concat / JaX / cpu / Primal 0.000003886798399980762 s 0.00000410532279997824 s 0.95
Concat / JaXPipe / cpu / Primal 0.0000039521924001746815 s 0.000004153241899985005 s 0.95
Concat / JaX / cpu / Forward 0.000005953432699971017 s 0.000007026319200031139 s 0.85
Concat / JaXPipe / cpu / Forward 0.000006147309600055451 s 0.000007335820499974943 s 0.84
Concat / JaX / cpu / BothRev 0.000006066617599935853 s 0.00000748516550002023 s 0.81
Concat / JaXPipe / cpu / PreRev 0.000006223077600225224 s 0.000007412402799991469 s 0.84
Concat / JaXPipe / cpu / PostRev 0.000006376959499903023 s 0.000007310956699984672 s 0.87
Concat / JaXPipe / cpu / BothRev 0.0000064608204997057326 s 0.000007627980400002343 s 0.85
Concat / JaX / tpu / Primal 0.0001521794316999 s 0.0001364663095009 s 1.12
Concat / JaXPipe / tpu / Primal 0.0001513736611002 s 0.0001385396195997 s 1.09
Concat / JaX / tpu / Forward 0.0002426793204998 s 0.0002357433504992 s 1.03
Concat / JaXPipe / tpu / Forward 0.0002391134818997 s 0.0002070721947995 s 1.15
Concat / JaX / tpu / BothRev 0.0002173405994002 s 0.000218590077701 s 0.99
Concat / JaXPipe / tpu / PreRev 0.0002252090182999 s 0.0002194721006002 s 1.03
Concat / JaXPipe / tpu / PostRev 0.0002363583399997 s 0.0002111700346998 s 1.12
Concat / JaXPipe / tpu / BothRev 0.0002351318683999 s 0.0002134689347003 s 1.10
Concat / JaX / cpu / Primal 0.000005281109700081288 s 0.00000410532279997824 s 1.29
Concat / JaXPipe / cpu / Primal 0.00000517447539996283 s 0.000004153241899985005 s 1.25
Concat / JaX / cpu / Forward 0.000008497035799973673 s 0.000007026319200031139 s 1.21
Concat / JaXPipe / cpu / Forward 0.000008991856600005122 s 0.000007335820499974943 s 1.23
Concat / JaX / cpu / BothRev 0.000008645671399972344 s 0.00000748516550002023 s 1.16
Concat / JaXPipe / cpu / PreRev 0.000009184412300055556 s 0.000007412402799991469 s 1.24
Concat / JaXPipe / cpu / PostRev 0.000008890117800001463 s 0.000007310956699984672 s 1.22
Concat / JaXPipe / cpu / BothRev 0.00000902029279995986 s 0.000007627980400002343 s 1.18
const_scatter / JaX / cpu / Primal 0.000011214400001335888 s 0.00000795500000094762 s 1.41
const_scatter / JaXPipe / cpu / Primal 0.000007201500011433382 s 0.000006355699997584452 s 1.13
const_scatter / JaX / cpu / Forward 0.000009912400037137558 s 0.000011652600005618296 s 0.85
const_scatter / JaXPipe / cpu / Forward 0.000010024800030805636 s 0.000013005499977225552 s 0.77
const_scatter / JaX / cpu / Primal 0.000012996300938539209 s 0.00000795500000094762 s 1.63
const_scatter / JaXPipe / cpu / Primal 0.000012446398613974452 s 0.000006355699997584452 s 1.96
const_scatter / JaX / cpu / Forward 0.00001559689990244806 s 0.000011652600005618296 s 1.34
const_scatter / JaXPipe / cpu / Forward 0.00001858269970398396 s 0.000013005499977225552 s 1.43
const_scatter / JaX / gpu / Primal 0.0001197963021695 s 0.0001228880999406 s 0.97
const_scatter / JaXPipe / gpu / Primal 0.0001087383017875 s 0.0001118395000958 s 0.97
const_scatter / JaX / gpu / Forward 0.0001451624004403 s 0.0001546240999232 s 0.94
const_scatter / JaXPipe / gpu / Forward 0.0001465364999603 s 0.0001445375999537 s 1.01
const_scatter / JaX / cpu / Primal 0.000006420900172088296 s 0.00000795500000094762 s 0.81
const_scatter / JaXPipe / cpu / Primal 0.000005788999987998977 s 0.000006355699997584452 s 0.91
const_scatter / JaX / cpu / Forward 0.000007339999865507707 s 0.000011652600005618296 s 0.63
const_scatter / JaXPipe / cpu / Forward 0.000008253000123659148 s 0.000013005499977225552 s 0.63
const_scatter / JaX / tpu / Primal 0.0001703129000816 s 0.0001590380008565 s 1.07
const_scatter / JaXPipe / tpu / Primal 0.0001711848999548 s 0.0001635230000829 s 1.05
const_scatter / JaX / tpu / Forward 0.0002528369001083 s 0.0002549369004555 s 0.99
const_scatter / JaXPipe / tpu / Forward 0.0002443608998873 s 0.000259205000475 s 0.94
const_scatter / JaX / cpu / Primal 0.000010082600056193767 s 0.00000795500000094762 s 1.27
const_scatter / JaXPipe / cpu / Primal 0.000008345100013684714 s 0.000006355699997584452 s 1.31
const_scatter / JaX / cpu / Forward 0.000014267100050346926 s 0.000011652600005618296 s 1.22
const_scatter / JaXPipe / cpu / Forward 0.000012778899963450383 s 0.000013005499977225552 s 0.98
GenDot / JaX / cpu / Primal 0.000004306294300022273 s 0.000004130862900001375 s 1.04
GenDot / JaXPipe / cpu / Primal 0.000004379963900009898 s 0.000004321130300013465 s 1.01
GenDot / JaX / cpu / Forward 0.0000064913046000128814 s 0.000006561183100029666 s 0.99
GenDot / JaXPipe / cpu / Forward 0.000007044539699973029 s 0.000007057765600029597 s 1.00
GenDot / JaX / cpu / BothRev 0.00000656351799998447 s 0.000006453430400006255 s 1.02
GenDot / JaXPipe / cpu / PreRev 0.000007100313099999766 s 0.000007266823100007969 s 0.98
GenDot / JaXPipe / cpu / PostRev 0.000006616024699997069 s 0.000006722887199975958 s 0.98
GenDot / JaXPipe / cpu / BothRev 0.0000071052775999760345 s 0.000007424744499985536 s 0.96
GenDot / JaX / cpu / Primal 0.000007744953100336715 s 0.000004130862900001375 s 1.87
GenDot / JaXPipe / cpu / Primal 0.000008510002301773057 s 0.000004321130300013465 s 1.97
GenDot / JaX / cpu / Forward 0.00001184704929764848 s 0.000006561183100029666 s 1.81
GenDot / JaXPipe / cpu / Forward 0.000014416873498703352 s 0.000007057765600029597 s 2.04
GenDot / JaX / cpu / BothRev 0.000011315846600336956 s 0.000006453430400006255 s 1.75
GenDot / JaXPipe / cpu / PreRev 0.000013962064098450354 s 0.000007266823100007969 s 1.92
GenDot / JaXPipe / cpu / PostRev 0.00001180869090021588 s 0.000006722887199975958 s 1.76
GenDot / JaXPipe / cpu / BothRev 0.00001551021019986365 s 0.000007424744499985536 s 2.09
GenDot / JaX / gpu / Primal 0.00007458653580106329 s 0.00007508556410011806 s 0.99
GenDot / JaXPipe / gpu / Primal 0.00008559405670093839 s 0.00008170207599996501 s 1.05
GenDot / JaX / gpu / Forward 0.0001019114563998 s 0.0001012580844 s 1.01
GenDot / JaXPipe / gpu / Forward 0.000132745992401 s 0.0001022659128 s 1.30
GenDot / JaX / gpu / BothRev 0.0001018532783025 s 0.000108609332 s 0.94
GenDot / JaXPipe / gpu / PreRev 0.0001039343890006 s 0.0001069992770999 s 0.97
GenDot / JaXPipe / gpu / PostRev 0.0001031373025994 s 0.0001041212174 s 0.99
GenDot / JaXPipe / gpu / BothRev 0.0001028495422011 s 0.0001014366046998 s 1.01
GenDot / JaX / cpu / Primal 0.000003967030400235671 s 0.000004130862900001375 s 0.96
GenDot / JaXPipe / cpu / Primal 0.000004261441399648902 s 0.000004321130300013465 s 0.99
GenDot / JaX / cpu / Forward 0.00000557618189995992 s 0.000006561183100029666 s 0.85
GenDot / JaXPipe / cpu / Forward 0.00000651672950007196 s 0.000007057765600029597 s 0.92
GenDot / JaX / cpu / BothRev 0.0000056783067997457694 s 0.000006453430400006255 s 0.88
GenDot / JaXPipe / cpu / PreRev 0.000006557558499844162 s 0.000007266823100007969 s 0.90
GenDot / JaXPipe / cpu / PostRev 0.000005962994700166746 s 0.000006722887199975958 s 0.89
GenDot / JaXPipe / cpu / BothRev 0.000006494498500251212 s 0.000007424744499985536 s 0.87
GenDot / JaX / tpu / Primal 0.0001357286365 s 0.0001500371514004 s 0.90
GenDot / JaXPipe / tpu / Primal 0.0001327944627999 s 0.0001476699334001 s 0.90
GenDot / JaX / tpu / Forward 0.0002317476936001 s 0.0002261418145004 s 1.02
GenDot / JaXPipe / tpu / Forward 0.0002304090810997 s 0.0002256598125997 s 1.02
GenDot / JaX / tpu / BothRev 0.0002242706475 s 0.0002336128614988 s 0.96
GenDot / JaXPipe / tpu / PreRev 0.0002389575138 s 0.0002078949758011 s 1.15
GenDot / JaXPipe / tpu / PostRev 0.0002304191911 s 0.0002292594575003 s 1.01
GenDot / JaXPipe / tpu / BothRev 0.0002409637979999 s 0.0002310455295009 s 1.04
GenDot / JaX / cpu / Primal 0.000005551975399976072 s 0.000004130862900001375 s 1.34
GenDot / JaXPipe / cpu / Primal 0.000005937156199979654 s 0.000004321130300013465 s 1.37
GenDot / JaX / cpu / Forward 0.000008200790099999723 s 0.000006561183100029666 s 1.25
GenDot / JaXPipe / cpu / Forward 0.000010054717900038669 s 0.000007057765600029597 s 1.42
GenDot / JaX / cpu / BothRev 0.000008273575400016852 s 0.000006453430400006255 s 1.28
GenDot / JaXPipe / cpu / PreRev 0.000009609804100000474 s 0.000007266823100007969 s 1.32
GenDot / JaXPipe / cpu / PostRev 0.000008123867900030745 s 0.000006722887199975958 s 1.21
GenDot / JaXPipe / cpu / BothRev 0.000009565223600020544 s 0.000007424744499985536 s 1.29
hlo_ffi / JaX / cpu / Primal 0.000005801384500000495 s 0.000006040187300004618 s 0.96
hlo_ffi / JaXPipe / cpu / Primal 0.00000585321019998446 s 0.000005982882199987216 s 0.98
hlo_ffi / JaX / cpu / Forward 0.000009361024299960263 s 0.000010758034999980735 s 0.87
hlo_ffi / JaXPipe / cpu / Forward 0.000009399197600032492 s 0.000010735428500038323 s 0.88
hlo_ffi / JaX / cpu / BothRev 0.000009114809800030344 s 0.00001012578710001435 s 0.90
hlo_ffi / JaXPipe / cpu / PreRev 0.000009112862099982522 s 0.00001022553389998393 s 0.89
hlo_ffi / JaXPipe / cpu / PostRev 0.000009296534800023436 s 0.000010337231099993005 s 0.90
hlo_ffi / JaXPipe / cpu / BothRev 0.000009130637900034345 s 0.000010255871799972735 s 0.89
hlo_ffi / JaX / cpu / Primal 0.000010959823301527648 s 0.000006040187300004618 s 1.81
hlo_ffi / JaXPipe / cpu / Primal 0.00001100334110087715 s 0.000005982882199987216 s 1.84
hlo_ffi / JaX / cpu / Forward 0.00001744639610114973 s 0.000010758034999980735 s 1.62
hlo_ffi / JaXPipe / cpu / Forward 0.000016957913598162122 s 0.000010735428500038323 s 1.58
hlo_ffi / JaX / cpu / BothRev 0.000016136582900071516 s 0.00001012578710001435 s 1.59
hlo_ffi / JaXPipe / cpu / PreRev 0.000016570974499336443 s 0.00001022553389998393 s 1.62
hlo_ffi / JaXPipe / cpu / PostRev 0.000017057824099902065 s 0.000010337231099993005 s 1.65
hlo_ffi / JaXPipe / cpu / BothRev 0.00001674101079988759 s 0.000010255871799972735 s 1.63
hlo_ffi / JaX / cpu / Primal 0.000005306748899965896 s 0.000006040187300004618 s 0.88
hlo_ffi / JaXPipe / cpu / Primal 0.000005064396099987789 s 0.000005982882199987216 s 0.85
hlo_ffi / JaX / cpu / Forward 0.000007620279000184383 s 0.000010758034999980735 s 0.71
hlo_ffi / JaXPipe / cpu / Forward 0.000007992232900141971 s 0.000010735428500038323 s 0.74
hlo_ffi / JaX / cpu / BothRev 0.000007733283000197844 s 0.00001012578710001435 s 0.76
hlo_ffi / JaXPipe / cpu / PreRev 0.000007728931000019656 s 0.00001022553389998393 s 0.76
hlo_ffi / JaXPipe / cpu / PostRev 0.000007702324999991105 s 0.000010337231099993005 s 0.75
hlo_ffi / JaXPipe / cpu / BothRev 0.000007752974999675644 s 0.000010255871799972735 s 0.76
hlo_ffi / JaX / tpu / Primal 0.00004894244690003689 s 0.00006669290239951806 s 0.73
hlo_ffi / JaXPipe / tpu / Primal 0.00006504267460004484 s 0.00006682222229865147 s 0.97
hlo_ffi / JaX / tpu / Forward 0.00009231488009972964 s 0.00009054184310080018 s 1.02
hlo_ffi / JaXPipe / tpu / Forward 0.0000937753824000538 s 0.0000919896151011926 s 1.02
hlo_ffi / JaX / tpu / BothRev 0.00009343231060011022 s 0.00009068617109878688 s 1.03
hlo_ffi / JaXPipe / tpu / PreRev 0.00009345764669997152 s 0.0000916849690998788 s 1.02
hlo_ffi / JaXPipe / tpu / PostRev 0.00009240192400029628 s 0.00009121268909948412 s 1.01
hlo_ffi / JaXPipe / tpu / BothRev 0.00009307556979983928 s 0.0000904399981998722 s 1.03
hlo_ffi / JaX / cpu / Primal 0.000007259182600046188 s 0.000006040187300004618 s 1.20
hlo_ffi / JaXPipe / cpu / Primal 0.00000709973800003354 s 0.000005982882199987216 s 1.19
hlo_ffi / JaX / cpu / Forward 0.000011007720099951258 s 0.000010758034999980735 s 1.02
hlo_ffi / JaXPipe / cpu / Forward 0.00001150235770001018 s 0.000010735428500038323 s 1.07
hlo_ffi / JaX / cpu / BothRev 0.0000107101161999708 s 0.00001012578710001435 s 1.06
hlo_ffi / JaXPipe / cpu / PreRev 0.000010969377099991109 s 0.00001022553389998393 s 1.07
hlo_ffi / JaXPipe / cpu / PostRev 0.00001108291949994964 s 0.000010337231099993005 s 1.07
hlo_ffi / JaXPipe / cpu / BothRev 0.000011339511200003472 s 0.000010255871799972735 s 1.11
llama / JaXPipe / cpu / Primal 0.0009074425700009 s 0.0008886572400024 s 1.02
llama / JaX / cpu / Primal 0.0009095653600024 s 0.0008988740899985 s 1.01
llama / HLOOpt / cpu / Primal 0.0009862448300009 s 0.0009684336400005 s 1.02
llama / PartOpt / cpu / Primal 0.0008789018399966 s 0.0008732787799999 s 1.01
llama / DefOpt / cpu / Primal 0.0009967112000003 s 0.0009657667599958 s 1.03
llama / IPartOpt / cpu / Primal 0.0008850172800021 s 0.0008895045300005 s 0.99
llama / IDefOpt / cpu / Primal 0.0009776081400013 s 0.0009823936600014 s 1.00
llama / JaXPipe / cpu / Forward 0.0026235451799993 s 0.0026352779900025 s 1.00
llama / JaX / cpu / Forward 0.0026558312799988 s 0.0026700162299994 s 0.99
llama / HLOOpt / cpu / Forward 0.002646916120002 s 0.0026948179600003 s 0.98
llama / PartOpt / cpu / Forward 0.00258393006 s 0.0026737331200001 s 0.97
llama / DefOpt / cpu / Forward 0.0026606134099984 s 0.0027148559099987 s 0.98
llama / IPartOpt / cpu / Forward 0.0026623046599979 s 0.0027034759200023 s 0.98
llama / IDefOpt / cpu / Forward 0.002609805889997 s 0.0027223944800016 s 0.96
llama / JaXPipe / cpu / PreRev 0.0024331153000002 s 0.0024359545000015 s 1.00
llama / JaXPipe / cpu / PostRev 0.0023608839999997 s 0.0025480002099993 s 0.93
llama / JaXPipe / cpu / BothRev 0.0023791065299974 s 0.002678289800001 s 0.89
llama / JaX / cpu / BothRev 0.0023183369400021 s 0.0024490471599983 s 0.95
llama / HLOOpt / cpu / PreRev 0.0028105332699988 s 0.0024723276899976 s 1.14
llama / HLOOpt / cpu / PostRev 0.002270092409999 s 0.0025622827000006 s 0.89
llama / HLOOpt / cpu / BothRev 0.0023200145200007 s 0.0025667688599969 s 0.90
llama / PartOpt / cpu / PreRev 0.0024737217399979 s 0.002715378390003 s 0.91
llama / PartOpt / cpu / PostRev 0.0024014230300008 s 0.0025989417899972 s 0.92
llama / PartOpt / cpu / BothRev 0.002457770459996 s 0.0025671774300008 s 0.96
llama / DefOpt / cpu / PreRev 0.0024431798600016 s 0.002672454970002 s 0.91
llama / DefOpt / cpu / PostRev 0.0022407618400029 s 0.0025091499600011 s 0.89
llama / DefOpt / cpu / BothRev 0.0024496573899978 s 0.0025303755200002 s 0.97
llama / IPartOpt / cpu / PreRev 0.0024221346500007 s 0.0026511285700007 s 0.91
llama / IPartOpt / cpu / PostRev 0.0023867917799998 s 0.002483636259999 s 0.96
llama / IPartOpt / cpu / BothRev 0.0027277661100015 s 0.0024590344600028 s 1.11
llama / IDefOpt / cpu / PreRev 0.0023291518799987 s 0.0026475106999987 s 0.88
llama / IDefOpt / cpu / PostRev 0.0026667156700023 s 0.0025744587500003 s 1.04
llama / IDefOpt / cpu / BothRev 0.0023556587800021 s 0.0026064670100004 s 0.90
llama / JaXPipe / gpu / Primal 0.0004577892979723 s 0.0004497240560012 s 1.02
llama / JaX / gpu / Primal 0.000439707989979 s 0.0004409631339985 s 1.00
llama / HLOOpt / gpu / Primal 0.0004288313899887 s 0.0004215804619998 s 1.02
llama / PartOpt / gpu / Primal 0.0004253768140333 s 0.0004225893479997 s 1.01
llama / DefOpt / gpu / Primal 0.0004545626040198 s 0.0004406088020004 s 1.03
llama / IPartOpt / gpu / Primal 0.0004290350800147 s 0.0004258870080011 s 1.01
llama / IDefOpt / gpu / Primal 0.0004398883599787 s 0.0004232599599999 s 1.04
llama / JaXPipe / gpu / Forward 0.0007271563920076 s 0.0007303671320005 s 1.00
llama / JaX / gpu / Forward 0.0007521811840124 s 0.000755226786001 s 1.00
llama / HLOOpt / gpu / Forward 0.0007285485100001 s 0.000721327444 s 1.01
llama / PartOpt / gpu / Forward 0.000732300943986 s 0.0007219144919981 s 1.01
llama / DefOpt / gpu / Forward 0.0007328091799863 s 0.0007399540980004 s 0.99
llama / IPartOpt / gpu / Forward 0.0007090260039549 s 0.0007615890859997 s 0.93
llama / IDefOpt / gpu / Forward 0.0007102053380222 s 0.0007173321160007 s 0.99
llama / JaXPipe / gpu / PreRev 0.0007682441640063 s 0.0008359755620003 s 0.92
llama / JaXPipe / gpu / PostRev 0.0007813354620011 s 0.000778637632 s 1.00
llama / JaXPipe / gpu / BothRev 0.0008291959260241 s 0.000823968648001 s 1.01
llama / JaX / gpu / BothRev 0.0008103626719675 s 0.0007995351699973 s 1.01
llama / HLOOpt / gpu / PreRev 0.0008153197920182 s 0.0007736044720004 s 1.05
llama / HLOOpt / gpu / PostRev 0.0008409873000346 s 0.0008369700279981 s 1.00
llama / HLOOpt / gpu / BothRev 0.0007738620979944 s 0.000789579780001 s 0.98
llama / PartOpt / gpu / PreRev 0.0007881506819976 s 0.0007680719800009 s 1.03
llama / PartOpt / gpu / PostRev 0.0008153438220033 s 0.0007905498 s 1.03
llama / PartOpt / gpu / BothRev 0.0008301956200157 s 0.0007959215419978 s 1.04
llama / DefOpt / gpu / PreRev 0.0007817672500386 s 0.0008118556420013 s 0.96
llama / DefOpt / gpu / PostRev 0.0007954919699695 s 0.0007940119640006 s 1.00
llama / DefOpt / gpu / BothRev 0.0007899293939699 s 0.000828213853998 s 0.95
llama / IPartOpt / gpu / PreRev 0.0007834862780291 s 0.000822383390001 s 0.95
llama / IPartOpt / gpu / PostRev 0.0007649889019667 s 0.0008178755299995 s 0.94
llama / IPartOpt / gpu / BothRev 0.0007727077240124 s 0.0007951335760008 s 0.97
llama / IDefOpt / gpu / PreRev 0.0007668799700331 s 0.0007986997199986 s 0.96
llama / IDefOpt / gpu / PostRev 0.0008101175380288 s 0.0007780378140014 s 1.04
llama / IDefOpt / gpu / BothRev 0.000823913900007 s 0.00082438043 s 1.00
llama / JaXPipe / tpu / Primal 0.0003854687499988 s 0.0003738564959785 s 1.03
llama / JaX / tpu / Primal 0.0003782905940024 s 0.0003541937760019 s 1.07
llama / HLOOpt / tpu / Primal 0.0003730469560032 s 0.0003426041179918 s 1.09
llama / PartOpt / tpu / Primal 0.0003861828719964 s 0.0003559281759953 s 1.09
llama / DefOpt / tpu / Primal 0.0003603251999957 s 0.0003352076559967 s 1.07
llama / IPartOpt / tpu / Primal 0.0003869348499938 s 0.0003563313159975 s 1.09
llama / IDefOpt / tpu / Primal 0.0003794858519977 s 0.0003448926560231 s 1.10
llama / JaXPipe / tpu / Forward 0.0005893325939978 s 0.0005368874759878 s 1.10
llama / JaX / tpu / Forward 0.0007185889020038 s 0.0006747420140018 s 1.06
llama / HLOOpt / tpu / Forward 0.0005859748539951 s 0.0005359630139719 s 1.09
llama / PartOpt / tpu / Forward 0.0005802676359962 s 0.0005374689540185 s 1.08
llama / DefOpt / tpu / Forward 0.0005860805740012 s 0.0005386405339813 s 1.09
llama / IPartOpt / tpu / Forward 0.0005810877960029 s 0.0005279682159889 s 1.10
llama / IDefOpt / tpu / Forward 0.000585689473999 s 0.0005322856940038 s 1.10
llama / JaXPipe / tpu / PreRev 0.0004077935219975 s 0.0004071840760007 s 1.00
llama / JaXPipe / tpu / PostRev 0.0003579273020004 s 0.0003571468559966 s 1.00
llama / JaXPipe / tpu / BothRev 0.0003918007259999 s 0.0003913024780049 s 1.00
llama / JaX / tpu / BothRev 0.0003575705600014 s 0.0003575916359841 s 1.00
llama / HLOOpt / tpu / PreRev 0.0003920678040012 s 0.000391396196006 s 1.00
llama / HLOOpt / tpu / PostRev 0.0003901341879973 s 0.000390183615993 s 1.00
llama / HLOOpt / tpu / BothRev 0.0003916821460006 s 0.0003916380959853 s 1.00
llama / PartOpt / tpu / PreRev 0.0003915000060005 s 0.0003916020159958 s 1.00
llama / PartOpt / tpu / PostRev 0.0003721759940017 s 0.0003719714359904 s 1.00
llama / PartOpt / tpu / BothRev 0.0003915840659974 s 0.000391423095978 s 1.00
llama / DefOpt / tpu / PreRev 0.0003916268859975 s 0.0003912544359918 s 1.00
llama / DefOpt / tpu / PostRev 0.0003830182480014 s 0.0003831635559909 s 1.00
llama / DefOpt / tpu / BothRev 0.0003917254460029 s 0.0003917173559893 s 1.00
llama / IPartOpt / tpu / PreRev 0.0003915724260004 s 0.0003916791160008 s 1.00
llama / IPartOpt / tpu / PostRev 0.0003717686939999 s 0.0003721168159972 s 1.00
llama / IPartOpt / tpu / BothRev 0.0003916158459978 s 0.0003913586359994 s 1.00
llama / IDefOpt / tpu / PreRev 0.0003918957260029 s 0.0003915143559861 s 1.00
llama / IDefOpt / tpu / PostRev 0.0003992906619969 s 0.0003993518560018 s 1.00
llama / IDefOpt / tpu / BothRev 0.000391655446001 s 0.0003911786959797 s 1.00
llama / JaXPipe / cpu / Primal 0.0024441745300009 s 0.0008886572400024 s 2.75
llama / JaX / cpu / Primal 0.0022567443699972 s 0.0008988740899985 s 2.51
llama / HLOOpt / cpu / Primal 0.0024169109099966 s 0.0009684336400005 s 2.50
llama / PartOpt / cpu / Primal 0.0023942625299969 s 0.0008732787799999 s 2.74
llama / DefOpt / cpu / Primal 0.0024256925999998 s 0.0009657667599958 s 2.51
llama / IPartOpt / cpu / Primal 0.0021259391400053 s 0.0008895045300005 s 2.39
llama / IDefOpt / cpu / Primal 0.0024190359200019 s 0.0009823936600014 s 2.46
llama / JaXPipe / cpu / Forward 0.0062637110299965 s 0.0026352779900025 s 2.38
llama / JaX / cpu / Forward 0.0066393355899981 s 0.0026700162299994 s 2.49
llama / HLOOpt / cpu / Forward 0.0065170172700072 s 0.0026948179600003 s 2.42
llama / PartOpt / cpu / Forward 0.006557281840005 s 0.0026737331200001 s 2.45
llama / DefOpt / cpu / Forward 0.0064207214999987 s 0.0027148559099987 s 2.37
llama / IPartOpt / cpu / Forward 0.0066701985700001 s 0.0027034759200023 s 2.47
llama / IDefOpt / cpu / Forward 0.0063020032799977 s 0.0027223944800016 s 2.31
llama / JaXPipe / cpu / PreRev 0.0056298741599948 s 0.0024359545000015 s 2.31
llama / JaXPipe / cpu / PostRev 0.0053546249800001 s 0.0025480002099993 s 2.10
llama / JaXPipe / cpu / BothRev 0.0055694356999993 s 0.002678289800001 s 2.08
llama / JaX / cpu / BothRev 0.005056876159997 s 0.0024490471599983 s 2.06
llama / HLOOpt / cpu / PreRev 0.0064372224999988 s 0.0024723276899976 s 2.60
llama / HLOOpt / cpu / PostRev 0.0052521482100019 s 0.0025622827000006 s 2.05
llama / HLOOpt / cpu / BothRev 0.0061724280399994 s 0.0025667688599969 s 2.40
llama / PartOpt / cpu / PreRev 0.0058125310499963 s 0.002715378390003 s 2.14
llama / PartOpt / cpu / PostRev 0.0056393627600027 s 0.0025989417899972 s 2.17
llama / PartOpt / cpu / BothRev 0.0058176545700007 s 0.0025671774300008 s 2.27
llama / DefOpt / cpu / PreRev 0.0061114042899953 s 0.002672454970002 s 2.29
llama / DefOpt / cpu / PostRev 0.0052143605499986 s 0.0025091499600011 s 2.08
llama / DefOpt / cpu / BothRev 0.0061252250500001 s 0.0025303755200002 s 2.42
llama / IPartOpt / cpu / PreRev 0.0057329246400058 s 0.0026511285700007 s 2.16
llama / IPartOpt / cpu / PostRev 0.0057084622100046 s 0.002483636259999 s 2.30
llama / IPartOpt / cpu / BothRev 0.0056965557000057 s 0.0024590344600028 s 2.32
llama / IDefOpt / cpu / PreRev 0.0061018636599965 s 0.0026475106999987 s 2.30
llama / IDefOpt / cpu / PostRev 0.0058290693700018 s 0.0025744587500003 s 2.26
llama / IDefOpt / cpu / BothRev 0.005656763070001 s 0.0026064670100004 s 2.17
scatter_sum / JaX / cpu / Primal 0.000004771006800001487 s 0.000004638723700008996 s 1.03
scatter_sum / JaXPipe / cpu / Primal 0.000004663005300017176 s 0.000004829928900016966 s 0.97
scatter_sum / JaX / cpu / Primal 0.000008974547099205665 s 0.000004638723700008996 s 1.93
scatter_sum / JaXPipe / cpu / Primal 0.000008781863600597717 s 0.000004829928900016966 s 1.82
scatter_sum / JaX / gpu / Primal 0.00008011844280117657 s 0.00008222425880012452 s 0.97
scatter_sum / JaXPipe / gpu / Primal 0.00007864698529883753 s 0.00008225327010004548 s 0.96
scatter_sum / JaX / cpu / Primal 0.00000453145229985239 s 0.000004638723700008996 s 0.98
scatter_sum / JaXPipe / cpu / Primal 0.000004514723299871548 s 0.000004829928900016966 s 0.93
scatter_sum / JaX / tpu / Primal 0.0001390461243001 s 0.0001515546452996 s 0.92
scatter_sum / JaXPipe / tpu / Primal 0.0001377100529 s 0.0001509173632992 s 0.91
scatter_sum / JaX / cpu / Primal 0.000006510114999946381 s 0.000004638723700008996 s 1.40
scatter_sum / JaXPipe / cpu / Primal 0.000006044843400013633 s 0.000004829928900016966 s 1.25
slicing / JaX / cpu / Primal 0.000003766015100018194 s 0.000003812621999986732 s 0.99
slicing / JaXPipe / cpu / Primal 0.000003680645399981586 s 0.000003800725999963106 s 0.97
slicing / JaX / cpu / Forward 0.000005817399500028841 s 0.000005971004199955132 s 0.97
slicing / JaXPipe / cpu / Forward 0.000005815487599966218 s 0.000005959274699989692 s 0.98
slicing / JaX / cpu / BothRev 0.000006307104699999399 s 0.000006395925000015268 s 0.99
slicing / JaXPipe / cpu / PreRev 0.000006302242900028432 s 0.000006324498000003587 s 1.00
slicing / JaXPipe / cpu / PostRev 0.000006341123099991819 s 0.000006327302800036705 s 1.00
slicing / JaXPipe / cpu / BothRev 0.000006284511900003054 s 0.000006399170399981813 s 0.98
slicing / JaX / cpu / Primal 0.000006818038001074456 s 0.000003812621999986732 s 1.79
slicing / JaXPipe / cpu / Primal 0.000007572723898920231 s 0.000003800725999963106 s 1.99
slicing / JaX / cpu / Forward 0.00001077530049951747 s 0.000005971004199955132 s 1.80
slicing / JaXPipe / cpu / Forward 0.000010678987900610082 s 0.000005959274699989692 s 1.79
slicing / JaX / cpu / BothRev 0.000011174885201035068 s 0.000006395925000015268 s 1.75
slicing / JaXPipe / cpu / PreRev 0.000012940334700397216 s 0.000006324498000003587 s 2.05
slicing / JaXPipe / cpu / PostRev 0.000011528543001622892 s 0.000006327302800036705 s 1.82
slicing / JaXPipe / cpu / BothRev 0.00001089338630263228 s 0.000006399170399981813 s 1.70
slicing / JaX / gpu / Primal 0.0000749777485994855 s 0.00008409172240008047 s 0.89
slicing / JaXPipe / gpu / Primal 0.00007469879339914769 s 0.00008141632549995847 s 0.92
slicing / JaX / gpu / Forward 0.0001167667011992 s 0.0001107257907999 s 1.05
slicing / JaXPipe / gpu / Forward 0.0001232135027996 s 0.0001054662010999 s 1.17
slicing / JaX / gpu / BothRev 0.0001030490562989 s 0.0001037999198 s 0.99
slicing / JaXPipe / gpu / PreRev 0.0001030302433005 s 0.0001054102685 s 0.98
slicing / JaXPipe / gpu / PostRev 0.0001009715944994 s 0.0001080322554 s 0.93
slicing / JaXPipe / gpu / BothRev 0.0001020368550991 s 0.0001041476307 s 0.98
slicing / JaX / cpu / Primal 0.000003762724600164802 s 0.000003812621999986732 s 0.99
slicing / JaXPipe / cpu / Primal 0.0000037691796002036423 s 0.000003800725999963106 s 0.99
slicing / JaX / cpu / Forward 0.000005004326099879108 s 0.000005971004199955132 s 0.84
slicing / JaXPipe / cpu / Forward 0.0000053861929998674895 s 0.000005959274699989692 s 0.90
slicing / JaX / cpu / BothRev 0.000005378232899965951 s 0.000006395925000015268 s 0.84
slicing / JaXPipe / cpu / PreRev 0.000005426598900157842 s 0.000006324498000003587 s 0.86
slicing / JaXPipe / cpu / PostRev 0.000005393267000181368 s 0.000006327302800036705 s 0.85
slicing / JaXPipe / cpu / BothRev 0.000005752482800016878 s 0.000006399170399981813 s 0.90
slicing / JaX / tpu / Primal 0.0001348638350002 s 0.0001530461573987 s 0.88
slicing / JaXPipe / tpu / Primal 0.0001527562231 s 0.0001470671273986 s 1.04
slicing / JaX / tpu / Forward 0.0002416819547997 s 0.0002250139306008 s 1.07
slicing / JaXPipe / tpu / Forward 0.0002546351777 s 0.0002034367818007 s 1.25
slicing / JaX / tpu / BothRev 0.0002576059055998 s 0.0001996357177995 s 1.29
slicing / JaXPipe / tpu / PreRev 0.0002472723934999 s 0.0002020191328992 s 1.22
slicing / JaXPipe / tpu / PostRev 0.0002303409371001 s 0.0002204992605998 s 1.04
slicing / JaXPipe / tpu / BothRev 0.0002275847081997 s 0.0002315618214997 s 0.98
slicing / JaX / cpu / Primal 0.0000047992813000746535 s 0.000003812621999986732 s 1.26
slicing / JaXPipe / cpu / Primal 0.000004806593200009957 s 0.000003800725999963106 s 1.26
slicing / JaX / cpu / Forward 0.000007643643799929123 s 0.000005971004199955132 s 1.28
slicing / JaXPipe / cpu / Forward 0.000007578653400014446 s 0.000005959274699989692 s 1.27
slicing / JaX / cpu / BothRev 0.000008075782000014443 s 0.000006395925000015268 s 1.26
slicing / JaXPipe / cpu / PreRev 0.000007651366900063294 s 0.000006324498000003587 s 1.21
slicing / JaXPipe / cpu / PostRev 0.000008070359300018026 s 0.000006327302800036705 s 1.28
slicing / JaXPipe / cpu / BothRev 0.000008082868900055473 s 0.000006399170399981813 s 1.26
sum / JaX / cpu / Primal 0.000004932808000012301 s 0.000004966364200026874 s 0.99
sum / JaXPipe / cpu / Primal 0.000005280869699981849 s 0.000004996704500035776 s 1.06
sum / JaX / cpu / Forward 0.000008105770900010612 s 0.000008281126400015636 s 0.98
sum / JaXPipe / cpu / Forward 0.000008148202499978653 s 0.000008336339800007408 s 0.98
sum / JaX / cpu / BothRev 0.000007425388000001476 s 0.000007623502599972199 s 0.97
sum / JaXPipe / cpu / PreRev 0.0000071321549999993295 s 0.000007313767599998755 s 0.98
sum / JaXPipe / cpu / PostRev 0.000007138050899993687 s 0.000007284650300016437 s 0.98
sum / JaXPipe / cpu / BothRev 0.000007206365100000767 s 0.000007277633999956379 s 0.99
sum / JaX / cpu / Primal 0.000010227256099460646 s 0.000004966364200026874 s 2.06
sum / JaXPipe / cpu / Primal 0.000010072741599287838 s 0.000004996704500035776 s 2.02
sum / JaX / cpu / Forward 0.000015981254298822022 s 0.000008281126400015636 s 1.93
sum / JaXPipe / cpu / Forward 0.000019085137100773864 s 0.000008336339800007408 s 2.29
sum / JaX / cpu / BothRev 0.000013984268801868891 s 0.000007623502599972199 s 1.83
sum / JaXPipe / cpu / PreRev 0.000013324989000102505 s 0.000007313767599998755 s 1.82
sum / JaXPipe / cpu / PostRev 0.00001403130700055044 s 0.000007284650300016437 s 1.93
sum / JaXPipe / cpu / BothRev 0.00001393604269833304 s 0.000007277633999956379 s 1.91
sum / JaX / gpu / Primal 0.00007295279990066775 s 0.00007549655999991956 s 0.97
sum / JaXPipe / gpu / Primal 0.00007293383380165324 s 0.00007316381340006046 s 1.00
sum / JaX / gpu / Forward 0.0001019623682019 s 0.00009866065129990602 s 1.03
sum / JaXPipe / gpu / Forward 0.000095873750600731 s 0.00009985945240005094 s 0.96
sum / JaX / gpu / BothRev 0.00009943614079966208 s 0.0001029208226 s 0.97
sum / JaXPipe / gpu / PreRev 0.0001066292684001 s 0.0001029314117 s 1.04
sum / JaXPipe / gpu / PostRev 0.0001178216299013 s 0.0001025825811999 s 1.15
sum / JaXPipe / gpu / BothRev 0.0001217962866998 s 0.0001006468552999 s 1.21
sum / JaX / cpu / Primal 0.000004528962300173589 s 0.000004966364200026874 s 0.91
sum / JaXPipe / cpu / Primal 0.000004614777199822129 s 0.000004996704500035776 s 0.92
sum / JaX / cpu / Forward 0.000007007190299918875 s 0.000008281126400015636 s 0.85
sum / JaXPipe / cpu / Forward 0.000007091206200129818 s 0.000008336339800007408 s 0.85
sum / JaX / cpu / BothRev 0.000006497161500010407 s 0.000007623502599972199 s 0.85
sum / JaXPipe / cpu / PreRev 0.000006258704599895282 s 0.000007313767599998755 s 0.86
sum / JaXPipe / cpu / PostRev 0.000006505797499994515 s 0.000007284650300016437 s 0.89
sum / JaXPipe / cpu / BothRev 0.000006494667500010109 s 0.000007277633999956379 s 0.89
sum / JaX / tpu / Primal 0.0001466714124002 s 0.0001627830832003 s 0.90
sum / JaXPipe / tpu / Primal 0.0001501904110999 s 0.0001509207273003 s 1.00
sum / JaX / tpu / Forward 0.0002219594902999 s 0.000217123673699 s 1.02
sum / JaXPipe / tpu / Forward 0.0002086299165002 s 0.0002189343637001 s 0.95
sum / JaX / tpu / BothRev 0.0002149565069998 s 0.0002243553595995 s 0.96
sum / JaXPipe / tpu / PreRev 0.0002389876208002 s 0.0002296060796012 s 1.04
sum / JaXPipe / tpu / PostRev 0.0002404344591999 s 0.0002212022575986 s 1.09
sum / JaXPipe / tpu / BothRev 0.0002283651327998 s 0.0002194539036994 s 1.04
sum / JaX / cpu / Primal 0.000006661810399964452 s 0.000004966364200026874 s 1.34
sum / JaXPipe / cpu / Primal 0.000006614429099954578 s 0.000004996704500035776 s 1.32
sum / JaX / cpu / Forward 0.000010513027100023464 s 0.000008281126400015636 s 1.27
sum / JaXPipe / cpu / Forward 0.000010886091500015029 s 0.000008336339800007408 s 1.31
sum / JaX / cpu / BothRev 0.000009496165400014431 s 0.000007623502599972199 s 1.25
sum / JaXPipe / cpu / PreRev 0.000009373181099999785 s 0.000007313767599998755 s 1.28
sum / JaXPipe / cpu / PostRev 0.000009255811799994262 s 0.000007284650300016437 s 1.27
sum / JaXPipe / cpu / BothRev 0.000009589105600025504 s 0.000007277633999956379 s 1.32
jaxmd40 / JaXPipe / gpu / Primal 0.001020465302281 s 0.0010153203000299 s 1.01
jaxmd40 / JaX / gpu / Primal 0.0010190708999289 s 0.0010191862000283 s 1.00
jaxmd40 / HLOOpt / gpu / Primal 0.0009678288013674 s 0.0009658715000114 s 1.00
jaxmd40 / PartOpt / gpu / Primal 0.0009576378011843 s 0.0009591229001671 s 1.00
jaxmd40 / DefOpt / gpu / Primal 0.0006880016007926 s 0.0006933840000783 s 0.99
jaxmd40 / IPartOpt / gpu / Primal 0.0009549135982524 s 0.000958172600076 s 1.00
jaxmd40 / IDefOpt / gpu / Primal 0.0007191376003902 s 0.0007199059000413 s 1.00
jaxmd40 / JaX / gpu / Forward 0.0012666150985751 s 0.0012737351999021 s 0.99
jaxmd40 / JaXPipe / gpu / PostRev 0.0040475852001691 s 0.0040386098999078 s 1.00
jaxmd40 / JaX / gpu / BothRev 0.0040457254013745 s 0.0040383210000072 s 1.00
jaxmd40 / HLOOpt / gpu / PostRev 0.0040508317993953 s 0.0040382227000009 s 1.00
jaxmd40 / PartOpt / gpu / PostRev 0.0041295258997706 s 0.0041242659999625 s 1.00
jaxmd40 / DefOpt / gpu / PostRev 0.0022101532988017 s 0.0022276448999036 s 0.99
jaxmd40 / IPartOpt / gpu / PostRev 0.0041285785991931 s 0.004175654500068 s 0.99
jaxmd40 / IDefOpt / gpu / PostRev 0.0026149545999942 s 0.0025700260999656 s 1.02
jaxmd40 / JaXPipe / tpu / Primal 0.00009070800006156788 s 0.00008584800089010969 s 1.06
jaxmd40 / JaX / tpu / Primal 0.00009838400001171976 s 0.00009776400111149997 s 1.01
jaxmd40 / HLOOpt / tpu / Primal 0.00009426200012967456 s 0.0001052790001267 s 0.90
jaxmd40 / PartOpt / tpu / Primal 0.0001003310000669 s 0.00009740399982547388 s 1.03
jaxmd40 / DefOpt / tpu / Primal 0.00009761789988260716 s 0.00009953199914889412 s 0.98
jaxmd40 / IPartOpt / tpu / Primal 0.0001018498998746 s 0.0001024729994242 s 0.99
jaxmd40 / IDefOpt / tpu / Primal 0.0001042830001097 s 0.0001011480009765 s 1.03
jaxmd40 / JaX / tpu / Forward 0.0001784639000106 s 0.0001878370007034 s 0.95
jaxmd40 / JaXPipe / tpu / PostRev 0.0001885019002656 s 0.0001794249998056 s 1.05
jaxmd40 / JaX / tpu / BothRev 0.0002088629000354 s 0.0002060239989077 s 1.01
jaxmd40 / HLOOpt / tpu / PostRev 0.0001984179001738 s 0.0002016290003666 s 0.98
jaxmd40 / PartOpt / tpu / PostRev 0.0002070499002002 s 0.0001994299993384 s 1.04
jaxmd40 / DefOpt / tpu / PostRev 0.0001958759003173 s 0.0001948909994098 s 1.01
jaxmd40 / IPartOpt / tpu / PostRev 0.0002073379000648 s 0.0001982829999178 s 1.05
jaxmd40 / IDefOpt / tpu / PostRev 0.0002032610002061 s 0.000202244000684 s 1.01
jaxmd40 / JaXPipe / cpu / Primal 0.00005737940000472008 s 0.00004885719999947469 s 1.17
jaxmd40 / JaX / cpu / Primal 0.00004496030005611829 s 0.00004359650001788396 s 1.03
jaxmd40 / HLOOpt / cpu / Primal 0.00006886660003146972 s 0.00006475049995060544 s 1.06
jaxmd40 / PartOpt / cpu / Primal 0.00006734499993399368 s 0.00006033259996911511 s 1.12
jaxmd40 / DefOpt / cpu / Primal 0.00006536219998451998 s 0.0000620943999820156 s 1.05
jaxmd40 / IPartOpt / cpu / Primal 0.00007061090000206604 s 0.00007177099996624747 s 0.98
jaxmd40 / IDefOpt / cpu / Primal 0.0000647492000098282 s 0.00006640310002694604 s 0.98
jaxmd40 / JaX / cpu / Forward 0.0001104166999539 s 0.0001079659000424 s 1.02
jaxmd40 / JaXPipe / cpu / PostRev 0.0001343590000033 s 0.000135948299976 s 0.99
jaxmd40 / JaX / cpu / BothRev 0.0001879305000329 s 0.0001765853000506 s 1.06
jaxmd40 / HLOOpt / cpu / PostRev 0.0001373711999804 s 0.0001395610000145 s 0.98
jaxmd40 / PartOpt / cpu / PostRev 0.000186272799965 s 0.0001681743999142 s 1.11
jaxmd40 / DefOpt / cpu / PostRev 0.0001316176000727 s 0.0001146505999713 s 1.15
jaxmd40 / IPartOpt / cpu / PostRev 0.000185376599984 s 0.0001595682000697 s 1.16
jaxmd40 / IDefOpt / cpu / PostRev 0.0001551435999317 s 0.000148431100024 s 1.05

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

@wsmoses
Copy link
Member

wsmoses commented Dec 18, 2025

actually I'm going to be rude and reject this PR as I realize we have equally copied the same function from polygeist into Enzyme proper: https://github.com/EnzymeAD/Enzyme/blob/0fdea7861949cdc9e2866e6b6f34c55fdcca2922/enzyme/Enzyme/MLIR/Interfaces/Utils.cpp#L136

Can we update the enzyme version, then remove the version here (and have it use the single version in enzyme proper)

@wsmoses
Copy link
Member

wsmoses commented Dec 18, 2025

though the fact it passes here regardless is a good sign

@ivanradanov ivanradanov changed the title Add support for LLVM noalias in alias analysis utils Use mlir utils from Enzyme Dec 18, 2025
@ivanradanov
Copy link
Collaborator Author

ivanradanov commented Dec 18, 2025

Moved here EnzymeAD/Enzyme#2635 and this PR now makes us use the Enzyme utils

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 Awaiting requested review from wsmoses

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