Skip to content

Add partial symmetry detection #1663

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
gaurav-arya wants to merge 27 commits into
base: main
Choose a base branch
from
Open

Add partial symmetry detection #1663

gaurav-arya wants to merge 27 commits into from

Conversation

@gaurav-arya
Copy link
Member

@gaurav-arya gaurav-arya commented Nov 28, 2025 •
edited
Loading

  • Implement partial symmetry lattice analysis
  • Port over + generalize detection logic from 2D symmetry analysis
  • Add initial optimization

@gaurav-arya gaurav-arya marked this pull request as ready for review November 28, 2025 07:51
github-actions[bot]
github-actions bot reviewed Nov 28, 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: f922615 Previous: 99d2b63 Ratio
actmtch / JaXPipe / cpu / Primal 0.000006623040007980308 s 0.000007074660015859991 s 0.94
actmtch / Jax / cpu / Primal 0.000006718440072290833 s 0.00000704749997566978 s 0.95
actmtch / HLOOpt / cpu / Primal 0.000007191480035544373 s 0.000007264960004249588 s 0.99
actmtch / PartOpt / cpu / Primal 0.000006342419965221779 s 0.00000608861999353394 s 1.04
actmtch / IPartOpt / cpu / Primal 0.000006567379969055764 s 0.000006828119985584636 s 0.96
actmtch / DefOpt / cpu / Primal 0.000007084820026648231 s 0.000008983800016721943 s 0.79
actmtch / IDefOpt / cpu / Primal 0.000007126260043150978 s 0.00000724573998923006 s 0.98
actmtch / JaXPipe / cpu / Forward 0.000011460839978099102 s 0.000011053500002162763 s 1.04
actmtch / Jax / cpu / Forward 0.000009248300084436778 s 0.000009771300019565388 s 0.95
actmtch / HLOOpt / cpu / Forward 0.000011537059981492348 s 0.000011151759999847854 s 1.03
actmtch / PartOpt / cpu / Forward 0.000010960959971271222 s 0.00001049553998200281 s 1.04
actmtch / IPartOpt / cpu / Forward 0.000011713519979821283 s 0.00001088953998987563 s 1.08
actmtch / DefOpt / cpu / Forward 0.00001053840007443796 s 0.000010273019979649687 s 1.03
actmtch / IDefOpt / cpu / Forward 0.000010812760083354078 s 0.000010934499978247913 s 0.99
actmtch / JaXPipe / cpu / PreRev 0.000011165539908688516 s 0.000011111580015494835 s 1.00
actmtch / JaXPipe / cpu / PostRev 0.000010262940031680046 s 0.000009981279990824987 s 1.03
actmtch / JaXPipe / cpu / BothRev 0.000010964459979732055 s 0.000011585639986151364 s 0.95
actmtch / Jax / cpu / BothRev 0.000009228020044247389 s 0.000009328440046374454 s 0.99
actmtch / HLOOpt / cpu / PreRev 0.000010851179922610754 s 0.000010915980010395289 s 0.99
actmtch / HLOOpt / cpu / PostRev 0.000012764159928337903 s 0.000012751240028592292 s 1.00
actmtch / HLOOpt / cpu / BothRev 0.00001097628000934492 s 0.000011103559972980291 s 0.99
actmtch / PartOpt / cpu / PreRev 0.000010606720079522348 s 0.00001094761998501781 s 0.97
actmtch / PartOpt / cpu / PostRev 0.0000097960799394059 s 0.000010008060016843955 s 0.98
actmtch / PartOpt / cpu / BothRev 0.000010756060000858269 s 0.000011785139986386638 s 0.91
actmtch / IPartOpt / cpu / PreRev 0.000010602580132399452 s 0.000010656700005711171 s 0.99
actmtch / IPartOpt / cpu / PostRev 0.00000963494003372034 s 0.000010013160017479094 s 0.96
actmtch / IPartOpt / cpu / BothRev 0.00001110526001866674 s 0.000011441900005593195 s 0.97
actmtch / DefOpt / cpu / PreRev 0.000011265360099059762 s 0.0000110185000266938 s 1.02
actmtch / DefOpt / cpu / PostRev 0.000010708459994930308 s 0.00001147651999417576 s 0.93
actmtch / DefOpt / cpu / BothRev 0.000010860859911190346 s 0.000010963859995172243 s 0.99
actmtch / IDefOpt / cpu / PreRev 0.000011273080053797455 s 0.0000109305400201265 s 1.03
actmtch / IDefOpt / cpu / PostRev 0.000010993720006808872 s 0.000011023619990737644 s 1.00
actmtch / IDefOpt / cpu / BothRev 0.000010757539967016782 s 0.000011036960022465792 s 0.97
actmtch / JaXPipe / cuda / Primal 0.000002016 s 0.000002047 s 0.98
actmtch / Jax / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / HLOOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / PartOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / IPartOpt / cuda / Primal 0.000002015 s 0.000002016 s 1.00
actmtch / DefOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / IDefOpt / cuda / Primal 0.000002016 s 0.000002016 s 1
actmtch / JaXPipe / cuda / Forward 0.00001136 s 0.00001008 s 1.13
actmtch / Jax / cuda / Forward 0.000010432 s 0.000010208 s 1.02
actmtch / HLOOpt / cuda / Forward 0.000009568 s 0.000009856 s 0.97
actmtch / PartOpt / cuda / Forward 0.000009568 s 0.000010305 s 0.93
actmtch / IPartOpt / cuda / Forward 0.000009792 s 0.000010176 s 0.96
actmtch / DefOpt / cuda / Forward 0.000009984 s 0.000013472 s 0.74
actmtch / IDefOpt / cuda / Forward 0.000009952 s 0.000010176 s 0.98
actmtch / JaXPipe / cuda / PreRev 0.000010111 s 0.000009856 s 1.03
actmtch / JaXPipe / cuda / PostRev 0.000009664 s 0.000010464 s 0.92
actmtch / JaXPipe / cuda / BothRev 0.00001008 s 0.000010047 s 1.00
actmtch / Jax / cuda / BothRev 0.000010368 s 0.000010016 s 1.04
actmtch / HLOOpt / cuda / PreRev 0.000010112 s 0.000009824 s 1.03
actmtch / HLOOpt / cuda / PostRev 0.000010016 s 0.00001024 s 0.98
actmtch / HLOOpt / cuda / BothRev 0.000009952 s 0.00001008 s 0.99
actmtch / PartOpt / cuda / PreRev 0.000009887 s 0.000010432 s 0.95
actmtch / PartOpt / cuda / PostRev 0.000010016 s 0.000010208 s 0.98
actmtch / PartOpt / cuda / BothRev 0.000010304 s 0.000010272 s 1.00
actmtch / IPartOpt / cuda / PreRev 0.000010336 s 0.000010144 s 1.02
actmtch / IPartOpt / cuda / PostRev 0.000010175 s 0.000010304 s 0.99
actmtch / IPartOpt / cuda / BothRev 0.000010144 s 0.000009792 s 1.04
actmtch / DefOpt / cuda / PreRev 0.000010145 s 0.000010176 s 1.00
actmtch / DefOpt / cuda / PostRev 0.000010176 s 0.000010048 s 1.01
actmtch / DefOpt / cuda / BothRev 0.000009824 s 0.000009728 s 1.01
actmtch / IDefOpt / cuda / PreRev 0.000009856 s 0.000010784 s 0.91
actmtch / IDefOpt / cuda / PostRev 0.000010048 s 0.000010303 s 0.98
actmtch / IDefOpt / cuda / BothRev 0.000010112 s 0.000010144 s 1.00
actmtch / JaXPipe / tpu / Primal 5.63425e-7 s 5.6375e-7 s 1.00
actmtch / Jax / tpu / Primal 5.965e-7 s 5.9715e-7 s 1.00
actmtch / HLOOpt / tpu / Primal 0.00000209465 s 0.000002100575 s 1.00
actmtch / PartOpt / tpu / Primal 5.9705e-7 s 5.965499999999999e-7 s 1.00
actmtch / IPartOpt / tpu / Primal 5.5255e-7 s 5.52725e-7 s 1.00
actmtch / DefOpt / tpu / Primal 0.0000021696 s 0.0000021609 s 1.00
actmtch / IDefOpt / tpu / Primal 0.00000211155 s 0.000002110575 s 1.00
actmtch / JaXPipe / tpu / Forward 0.0000038386 s 0.000003824075 s 1.00
actmtch / Jax / tpu / Forward 0.00000120895 s 0.00000121035 s 1.00
actmtch / HLOOpt / tpu / Forward 0.000003951525 s 0.000003934275 s 1.00
actmtch / PartOpt / tpu / Forward 0.000003911 s 0.000003911724999999999 s 1.00
actmtch / IPartOpt / tpu / Forward 0.00000393995 s 0.000003933925 s 1.00
actmtch / DefOpt / tpu / Forward 0.000003913499999999999 s 0.000003913074999999999 s 1.00
actmtch / IDefOpt / tpu / Forward 0.000003936824999999999 s 0.0000039450750000000005 s 1.00
actmtch / JaXPipe / tpu / PreRev 0.0000034844000000000005 s 0.0000034795 s 1.00
actmtch / JaXPipe / tpu / PostRev 0.00000164005 s 0.000001645175 s 1.00
actmtch / JaXPipe / tpu / BothRev 0.000003483 s 0.000003478825 s 1.00
actmtch / Jax / tpu / BothRev 0.0000016322 s 0.000001631225 s 1.00
actmtch / HLOOpt / tpu / PreRev 0.000003487075 s 0.0000034729500000000004 s 1.00
actmtch / HLOOpt / tpu / PostRev 0.000003410875 s 0.000003397575 s 1.00
actmtch / HLOOpt / tpu / BothRev 0.000003470375 s 0.000003481925 s 1.00
actmtch / PartOpt / tpu / PreRev 0.0000034137 s 0.0000034036 s 1.00
actmtch / PartOpt / tpu / PostRev 0.0000015908 s 0.000001587925 s 1.00
actmtch / PartOpt / tpu / BothRev 0.000003422125 s 0.000003399525 s 1.01
actmtch / IPartOpt / tpu / PreRev 0.000003472825 s 0.000003505025 s 0.99
actmtch / IPartOpt / tpu / PostRev 0.00000163855 s 0.000001635125 s 1.00
actmtch / IPartOpt / tpu / BothRev 0.000003489275 s 0.00000346705 s 1.01
actmtch / DefOpt / tpu / PreRev 0.0000034106 s 0.0000034026499999999995 s 1.00
actmtch / DefOpt / tpu / PostRev 0.000003415025 s 0.0000034195 s 1.00
actmtch / DefOpt / tpu / BothRev 0.000003416475 s 0.000003396175 s 1.01
actmtch / IDefOpt / tpu / PreRev 0.000003462325 s 0.00000349265 s 0.99
actmtch / IDefOpt / tpu / PostRev 0.000003404025 s 0.00000341735 s 1.00
actmtch / IDefOpt / tpu / BothRev 0.000003467450000000001 s 0.000003478575 s 1.00
actmtch / JaXPipe / cpu / Primal 0.000013349 s 0.000007074660015859991 s 1.89
actmtch / Jax / cpu / Primal 0.000013239 s 0.00000704749997566978 s 1.88
actmtch / HLOOpt / cpu / Primal 0.000013932 s 0.000007264960004249588 s 1.92
actmtch / PartOpt / cpu / Primal 0.000013399 s 0.00000608861999353394 s 2.20
actmtch / IPartOpt / cpu / Primal 0.000013146 s 0.000006828119985584636 s 1.93
actmtch / DefOpt / cpu / Primal 0.000014225 s 0.000008983800016721943 s 1.58
actmtch / IDefOpt / cpu / Primal 0.000014068 s 0.00000724573998923006 s 1.94
actmtch / JaXPipe / cpu / Forward 0.000019313 s 0.000011053500002162763 s 1.75
actmtch / Jax / cpu / Forward 0.000017977 s 0.000009771300019565388 s 1.84
actmtch / HLOOpt / cpu / Forward 0.000019327 s 0.000011151759999847854 s 1.73
actmtch / PartOpt / cpu / Forward 0.000019276000000000003 s 0.00001049553998200281 s 1.84
actmtch / IPartOpt / cpu / Forward 0.000019474 s 0.00001088953998987563 s 1.79
actmtch / DefOpt / cpu / Forward 0.000019024 s 0.000010273019979649687 s 1.85
actmtch / IDefOpt / cpu / Forward 0.000018889 s 0.000010934499978247913 s 1.73
actmtch / JaXPipe / cpu / PreRev 0.000019343 s 0.000011111580015494835 s 1.74
actmtch / JaXPipe / cpu / PostRev 0.000017665 s 0.000009981279990824987 s 1.77
actmtch / JaXPipe / cpu / BothRev 0.000019305 s 0.000011585639986151364 s 1.67
actmtch / Jax / cpu / BothRev 0.000017859 s 0.000009328440046374454 s 1.91
actmtch / HLOOpt / cpu / PreRev 0.00001945 s 0.000010915980010395289 s 1.78
actmtch / HLOOpt / cpu / PostRev 0.000031163 s 0.000012751240028592292 s 2.44
actmtch / HLOOpt / cpu / BothRev 0.000019464 s 0.000011103559972980291 s 1.75
actmtch / PartOpt / cpu / PreRev 0.000019318 s 0.00001094761998501781 s 1.76
actmtch / PartOpt / cpu / PostRev 0.000017704999999999997 s 0.000010008060016843955 s 1.77
actmtch / PartOpt / cpu / BothRev 0.000019555 s 0.000011785139986386638 s 1.66
actmtch / IPartOpt / cpu / PreRev 0.000019292 s 0.000010656700005711171 s 1.81
actmtch / IPartOpt / cpu / PostRev 0.000017662 s 0.000010013160017479094 s 1.76
actmtch / IPartOpt / cpu / BothRev 0.000019676 s 0.000011441900005593195 s 1.72
actmtch / DefOpt / cpu / PreRev 0.000019175 s 0.0000110185000266938 s 1.74
actmtch / DefOpt / cpu / PostRev 0.00001936 s 0.00001147651999417576 s 1.69
actmtch / DefOpt / cpu / BothRev 0.000019333 s 0.000010963859995172243 s 1.76
actmtch / IDefOpt / cpu / PreRev 0.000019301 s 0.0000109305400201265 s 1.77
actmtch / IDefOpt / cpu / PostRev 0.00001914 s 0.000011023619990737644 s 1.74
actmtch / IDefOpt / cpu / BothRev 0.000019144 s 0.000011036960022465792 s 1.73
add_one / JaXPipe / cpu / Primal 0.000006525179942400427 s 0.000006555000009029755 s 1.00
add_one / Jax / cpu / Primal 0.000006542260071000782 s 0.000008358779950867756 s 0.78
add_one / HLOOpt / cpu / Primal 0.0000064999599999282506 s 0.000007089640039339429 s 0.92
add_one / PartOpt / cpu / Primal 0.000006332800021482399 s 0.00000640845996713324 s 0.99
add_one / IPartOpt / cpu / Primal 0.000006833360093878582 s 0.000006578960001206724 s 1.04
add_one / DefOpt / cpu / Primal 0.000006387320026988163 s 0.000006600820006497088 s 0.97
add_one / IDefOpt / cpu / Primal 0.000006526499983010581 s 0.000006797320029363618 s 0.96
add_one / JaXPipe / cpu / Forward 0.000009546280089125504 s 0.000009917019997374155 s 0.96
add_one / Jax / cpu / Forward 0.00000947533988437499 s 0.000009898720018099992 s 0.96
add_one / HLOOpt / cpu / Forward 0.000010447959994053235 s 0.0000102036600401334 s 1.02
add_one / PartOpt / cpu / Forward 0.00001000209997073398 s 0.000010513300003367476 s 0.95
add_one / IPartOpt / cpu / Forward 0.000009972339939849916 s 0.000010385160030637054 s 0.96
add_one / DefOpt / cpu / Forward 0.00000948200018683565 s 0.000010184639986619004 s 0.93
add_one / IDefOpt / cpu / Forward 0.000009636459963076047 s 0.000010276100047121872 s 0.94
add_one / JaXPipe / cpu / PreRev 0.000011080600033892552 s 0.000011770159999286988 s 0.94
add_one / JaXPipe / cpu / PostRev 0.000011664160047075713 s 0.000011907140014955076 s 0.98
add_one / JaXPipe / cpu / BothRev 0.000011809580046246992 s 0.000012559820015667356 s 0.94
add_one / Jax / cpu / BothRev 0.0000114413199480623 s 0.00001256429996828956 s 0.91
add_one / HLOOpt / cpu / PreRev 0.000011401839929021662 s 0.000012178099977973031 s 0.94
add_one / HLOOpt / cpu / PostRev 0.000016792039987194584 s 0.00001428686002327595 s 1.18
add_one / HLOOpt / cpu / BothRev 0.000011853080068249257 s 0.00001226867996592773 s 0.97
add_one / PartOpt / cpu / PreRev 0.000010965720011881786 s 0.00001211154001794057 s 0.91
add_one / PartOpt / cpu / PostRev 0.000011234359953959938 s 0.000012167160066383077 s 0.92
add_one / PartOpt / cpu / BothRev 0.000011719439917214914 s 0.000012089080009900498 s 0.97
add_one / IPartOpt / cpu / PreRev 0.000011175220079167047 s 0.000012506240018410608 s 0.89
add_one / IPartOpt / cpu / PostRev 0.000011697479967551772 s 0.000011844299979202334 s 0.99
add_one / IPartOpt / cpu / BothRev 0.000011066000024584354 s 0.00001200290000269888 s 0.92
add_one / DefOpt / cpu / PreRev 0.000011083499994128942 s 0.000012106419962947256 s 0.92
add_one / DefOpt / cpu / PostRev 0.00001158920000307262 s 0.000011705979941325495 s 0.99
add_one / DefOpt / cpu / BothRev 0.00001148910001575132 s 0.00001240259999576665 s 0.93
add_one / IDefOpt / cpu / PreRev 0.00001084837997041177 s 0.00001211826000144356 s 0.90
add_one / IDefOpt / cpu / PostRev 0.000011281619990768376 s 0.000012116420002712402 s 0.93
add_one / IDefOpt / cpu / BothRev 0.000011287680008535972 s 0.000011627660005615326 s 0.97
add_one / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / Jax / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / IPartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_one / JaXPipe / cuda / Forward 0.000010144 s 0.00001024 s 0.99
add_one / Jax / cuda / Forward 0.00000928 s 0.000010368 s 0.90
add_one / HLOOpt / cuda / Forward 0.000010208 s 0.000010208 s 1
add_one / PartOpt / cuda / Forward 0.000011008 s 0.00001024 s 1.07
add_one / IPartOpt / cuda / Forward 0.000009984 s 0.000010464 s 0.95
add_one / DefOpt / cuda / Forward 0.000011104 s 0.000010208 s 1.09
add_one / IDefOpt / cuda / Forward 0.000010176 s 0.000010304 s 0.99
add_one / JaXPipe / cuda / PreRev 0.000024416 s 0.000025728 s 0.95
add_one / JaXPipe / cuda / PostRev 0.00002432 s 0.000025536 s 0.95
add_one / JaXPipe / cuda / BothRev 0.00002432 s 0.000024512 s 0.99
add_one / Jax / cuda / BothRev 0.000024384 s 0.000025568 s 0.95
add_one / HLOOpt / cuda / PreRev 0.000024991 s 0.000025728 s 0.97
add_one / HLOOpt / cuda / PostRev 0.00002736 s 0.00002528 s 1.08
add_one / HLOOpt / cuda / BothRev 0.000024096 s 0.000024737 s 0.97
add_one / PartOpt / cuda / PreRev 0.000024832 s 0.000025152 s 0.99
add_one / PartOpt / cuda / PostRev 0.000024895 s 0.000029312 s 0.85
add_one / PartOpt / cuda / BothRev 0.000024608 s 0.000029568 s 0.83
add_one / IPartOpt / cuda / PreRev 0.000024672 s 0.000029248 s 0.84
add_one / IPartOpt / cuda / PostRev 0.000024544 s 0.000024609 s 1.00
add_one / IPartOpt / cuda / BothRev 0.000024384 s 0.00002448 s 1.00
add_one / DefOpt / cuda / PreRev 0.000024928 s 0.000025377 s 0.98
add_one / DefOpt / cuda / PostRev 0.000025024 s 0.000025152 s 0.99
add_one / DefOpt / cuda / BothRev 0.000024449 s 0.000024736 s 0.99
add_one / IDefOpt / cuda / PreRev 0.000024513 s 0.000026048 s 0.94
add_one / IDefOpt / cuda / PostRev 0.000024 s 0.000025536 s 0.94
add_one / IDefOpt / cuda / BothRev 0.000024256 s 0.000025856 s 0.94
add_one / JaXPipe / tpu / Primal 0.0000014223250000000002 s 0.0000014431 s 0.99
add_one / Jax / tpu / Primal 0.0000014087 s 0.00000141315 s 1.00
add_one / HLOOpt / tpu / Primal 0.000001428 s 0.000001427425 s 1.00
add_one / PartOpt / tpu / Primal 0.0000014092750000000002 s 0.0000014064750000000002 s 1.00
add_one / IPartOpt / tpu / Primal 0.000001424575 s 0.000001427325 s 1.00
add_one / DefOpt / tpu / Primal 0.0000014008999999999998 s 0.000001403125 s 1.00
add_one / IDefOpt / tpu / Primal 0.0000014235999999999998 s 0.0000014284 s 1.00
add_one / JaXPipe / tpu / Forward 0.0000018434 s 0.0000018557 s 0.99
add_one / Jax / tpu / Forward 0.00000184555 s 0.000001852975 s 1.00
add_one / HLOOpt / tpu / Forward 0.000001859325 s 0.000001859325 s 1
add_one / PartOpt / tpu / Forward 0.00000183755 s 0.000001843425 s 1.00
add_one / IPartOpt / tpu / Forward 0.00000186115 s 0.000001849025 s 1.01
add_one / DefOpt / tpu / Forward 0.0000018409 s 0.000001850225 s 0.99
add_one / IDefOpt / tpu / Forward 0.0000018539 s 0.000001850875 s 1.00
add_one / JaXPipe / tpu / PreRev 0.0000022407 s 0.000002249275 s 1.00
add_one / JaXPipe / tpu / PostRev 0.0000022396 s 0.000002244275 s 1.00
add_one / JaXPipe / tpu / BothRev 0.0000022385 s 0.0000022332 s 1.00
add_one / Jax / tpu / BothRev 0.0000022433 s 0.0000022404 s 1.00
add_one / HLOOpt / tpu / PreRev 0.00000223915 s 0.000002238025 s 1.00
add_one / HLOOpt / tpu / PostRev 0.00000224085 s 0.00000224105 s 1.00
add_one / HLOOpt / tpu / BothRev 0.0000022323 s 0.0000022355 s 1.00
add_one / PartOpt / tpu / PreRev 0.00000223895 s 0.000002233825 s 1.00
add_one / PartOpt / tpu / PostRev 0.000002234425 s 0.000002244075 s 1.00
add_one / PartOpt / tpu / BothRev 0.0000022488 s 0.0000022386 s 1.00
add_one / IPartOpt / tpu / PreRev 0.0000022296500000000003 s 0.0000022358 s 1.00
add_one / IPartOpt / tpu / PostRev 0.000002248025 s 0.0000022391 s 1.00
add_one / IPartOpt / tpu / BothRev 0.000002239 s 0.000002234875 s 1.00
add_one / DefOpt / tpu / PreRev 0.00000224865 s 0.0000022432 s 1.00
add_one / DefOpt / tpu / PostRev 0.000002237625 s 0.0000022361 s 1.00
add_one / DefOpt / tpu / BothRev 0.000002236425 s 0.00000224885 s 0.99
add_one / IDefOpt / tpu / PreRev 0.0000022343 s 0.000002238725 s 1.00
add_one / IDefOpt / tpu / PostRev 0.00000223885 s 0.000002251525 s 0.99
add_one / IDefOpt / tpu / BothRev 0.000002245925 s 0.000002251525 s 1.00
add_one / JaXPipe / cpu / Primal 0.000012592 s 0.000006555000009029755 s 1.92
add_one / Jax / cpu / Primal 0.000013168 s 0.000008358779950867756 s 1.58
add_one / HLOOpt / cpu / Primal 0.000012841 s 0.000007089640039339429 s 1.81
add_one / PartOpt / cpu / Primal 0.000012835 s 0.00000640845996713324 s 2.00
add_one / IPartOpt / cpu / Primal 0.000013133 s 0.000006578960001206724 s 2.00
add_one / DefOpt / cpu / Primal 0.000013016 s 0.000006600820006497088 s 1.97
add_one / IDefOpt / cpu / Primal 0.000012809 s 0.000006797320029363618 s 1.88
add_one / JaXPipe / cpu / Forward 0.000017878 s 0.000009917019997374155 s 1.80
add_one / Jax / cpu / Forward 0.000017695 s 0.000009898720018099992 s 1.79
add_one / HLOOpt / cpu / Forward 0.000017573 s 0.0000102036600401334 s 1.72
add_one / PartOpt / cpu / Forward 0.000017767 s 0.000010513300003367476 s 1.69
add_one / IPartOpt / cpu / Forward 0.00001791 s 0.000010385160030637054 s 1.72
add_one / DefOpt / cpu / Forward 0.000017841 s 0.000010184639986619004 s 1.75
add_one / IDefOpt / cpu / Forward 0.00001798 s 0.000010276100047121872 s 1.75
add_one / JaXPipe / cpu / PreRev 0.000029861 s 0.000011770159999286988 s 2.54
add_one / JaXPipe / cpu / PostRev 0.000019981 s 0.000011907140014955076 s 1.68
add_one / JaXPipe / cpu / BothRev 0.000020287 s 0.000012559820015667356 s 1.62
add_one / Jax / cpu / BothRev 0.000019367 s 0.00001256429996828956 s 1.54
add_one / HLOOpt / cpu / PreRev 0.000019701 s 0.000012178099977973031 s 1.62
add_one / HLOOpt / cpu / PostRev 0.000019945 s 0.00001428686002327595 s 1.40
add_one / HLOOpt / cpu / BothRev 0.000019682 s 0.00001226867996592773 s 1.60
add_one / PartOpt / cpu / PreRev 0.000019136 s 0.00001211154001794057 s 1.58
add_one / PartOpt / cpu / PostRev 0.000019539 s 0.000012167160066383077 s 1.61
add_one / PartOpt / cpu / BothRev 0.000019599 s 0.000012089080009900498 s 1.62
add_one / IPartOpt / cpu / PreRev 0.000019831 s 0.000012506240018410608 s 1.59
add_one / IPartOpt / cpu / PostRev 0.000019883000000000003 s 0.000011844299979202334 s 1.68
add_one / IPartOpt / cpu / BothRev 0.000019307 s 0.00001200290000269888 s 1.61
add_one / DefOpt / cpu / PreRev 0.000019381 s 0.000012106419962947256 s 1.60
add_one / DefOpt / cpu / PostRev 0.000019632 s 0.000011705979941325495 s 1.68
add_one / DefOpt / cpu / BothRev 0.000019771 s 0.00001240259999576665 s 1.59
add_one / IDefOpt / cpu / PreRev 0.000019627 s 0.00001211826000144356 s 1.62
add_one / IDefOpt / cpu / PostRev 0.000020143 s 0.000012116420002712402 s 1.66
add_one / IDefOpt / cpu / BothRev 0.000019822 s 0.000011627660005615326 s 1.70
add_two / JaXPipe / cpu / Primal 0.000007139479930629023 s 0.000006799080001655966 s 1.05
add_two / Jax / cpu / Primal 0.000006654299922956852 s 0.000007222659960461897 s 0.92
add_two / HLOOpt / cpu / Primal 0.000007498019967897562 s 0.000006997399977990426 s 1.07
add_two / PartOpt / cpu / Primal 0.000006551419992320007 s 0.000006764840018149698 s 0.97
add_two / IPartOpt / cpu / Primal 0.000007541339946328662 s 0.000007197699969765381 s 1.05
add_two / DefOpt / cpu / Primal 0.000006857100106572034 s 0.00000721067998711078 s 0.95
add_two / IDefOpt / cpu / Primal 0.000006724160066369222 s 0.000006686059987259796 s 1.01
add_two / JaXPipe / cpu / Forward 0.000009919019994413248 s 0.00000990931998785527 s 1.00
add_two / Jax / cpu / Forward 0.000010289119891240262 s 0.000010424320007587083 s 0.99
add_two / HLOOpt / cpu / Forward 0.000010139559908566298 s 0.000010333659993193578 s 0.98
add_two / PartOpt / cpu / Forward 0.000009923779998644022 s 0.000010246299971186093 s 0.97
add_two / IPartOpt / cpu / Forward 0.00001051564000590588 s 0.000010061119974125175 s 1.05
add_two / DefOpt / cpu / Forward 0.00000987641997198807 s 0.00001001847997031291 s 0.99
add_two / IDefOpt / cpu / Forward 0.000010078440063807649 s 0.000010443999999552032 s 0.96
add_two / JaXPipe / cpu / PreRev 0.000014076819988986244 s 0.000014581140003429028 s 0.97
add_two / JaXPipe / cpu / PostRev 0.00001418780006133602 s 0.000014217040024959715 s 1.00
add_two / JaXPipe / cpu / BothRev 0.00001394950004396378 s 0.000014091619987084412 s 0.99
add_two / Jax / cpu / BothRev 0.00001421901997673558 s 0.000014226019975467351 s 1.00
add_two / HLOOpt / cpu / PreRev 0.000013925520033808424 s 0.0000144150200412696 s 0.97
add_two / HLOOpt / cpu / PostRev 0.000016253199992206645 s 0.000018896840019806403 s 0.86
add_two / HLOOpt / cpu / BothRev 0.000013753859911957987 s 0.000013867099996787146 s 0.99
add_two / PartOpt / cpu / PreRev 0.000014311400009319189 s 0.000014706779975313113 s 0.97
add_two / PartOpt / cpu / PostRev 0.000013572959942393936 s 0.000014247260014599306 s 0.95
add_two / PartOpt / cpu / BothRev 0.000014081059944146546 s 0.000014383599991560911 s 0.98
add_two / IPartOpt / cpu / PreRev 0.000014062640038901007 s 0.00001433168001312879 s 0.98
add_two / IPartOpt / cpu / PostRev 0.000013846920010109898 s 0.000014152100011415314 s 0.98
add_two / IPartOpt / cpu / BothRev 0.000014455359942076027 s 0.000014447140029005822 s 1.00
add_two / DefOpt / cpu / PreRev 0.000013582280043920036 s 0.000014267520018620417 s 0.95
add_two / DefOpt / cpu / PostRev 0.00001372591994368122 s 0.000014259899999160552 s 0.96
add_two / DefOpt / cpu / BothRev 0.00001386627998726908 s 0.000014279120023275029 s 0.97
add_two / IDefOpt / cpu / PreRev 0.00001385378005579696 s 0.00001455288002944144 s 0.95
add_two / IDefOpt / cpu / PostRev 0.000014118920007604177 s 0.000014553979999618603 s 0.97
add_two / IDefOpt / cpu / BothRev 0.000013777159965684405 s 0.000014658500022051158 s 0.94
add_two / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / Jax / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / IPartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
add_two / JaXPipe / cuda / Forward 0.000009888 s 0.00001008 s 0.98
add_two / Jax / cuda / Forward 0.00000992 s 0.000009824 s 1.01
add_two / HLOOpt / cuda / Forward 0.00000976 s 0.000010176 s 0.96
add_two / PartOpt / cuda / Forward 0.00000992 s 0.000009888 s 1.00
add_two / IPartOpt / cuda / Forward 0.000009473 s 0.00000944 s 1.00
add_two / DefOpt / cuda / Forward 0.000009504 s 0.000009824 s 0.97
add_two / IDefOpt / cuda / Forward 0.000009695 s 0.000010176 s 0.95
add_two / JaXPipe / cuda / PreRev 0.000032160000000000004 s 0.000032449 s 0.99
add_two / JaXPipe / cuda / PostRev 0.000032513 s 0.00003296 s 0.99
add_two / JaXPipe / cuda / BothRev 0.000032127999999999995 s 0.000032832 s 0.98
add_two / Jax / cuda / BothRev 0.000031392 s 0.000032416 s 0.97
add_two / HLOOpt / cuda / PreRev 0.000031936 s 0.000032767999999999995 s 0.97
add_two / HLOOpt / cuda / PostRev 0.000031328 s 0.000031777 s 0.99
add_two / HLOOpt / cuda / BothRev 0.000031296 s 0.000032672 s 0.96
add_two / PartOpt / cuda / PreRev 0.00003168 s 0.000032256 s 0.98
add_two / PartOpt / cuda / PostRev 0.000031712 s 0.000032160000000000004 s 0.99
add_two / PartOpt / cuda / BothRev 0.000032160000000000004 s 0.000031488 s 1.02
add_two / IPartOpt / cuda / PreRev 0.000031904000000000005 s 0.000032447 s 0.98
add_two / IPartOpt / cuda / PostRev 0.000032416 s 0.000032064 s 1.01
add_two / IPartOpt / cuda / BothRev 0.000031264 s 0.000032832 s 0.95
add_two / DefOpt / cuda / PreRev 0.000032127999999999995 s 0.000032448 s 0.99
add_two / DefOpt / cuda / PostRev 0.000031424 s 0.000032928 s 0.95
add_two / DefOpt / cuda / BothRev 0.000031807 s 0.000032672 s 0.97
add_two / IDefOpt / cuda / PreRev 0.000031936 s 0.000032384 s 0.99
add_two / IDefOpt / cuda / PostRev 0.000031711 s 0.00003296 s 0.96
add_two / IDefOpt / cuda / BothRev 0.000031392 s 0.000032 s 0.98
add_two / JaXPipe / tpu / Primal 0.00000143315 s 0.000001439975 s 1.00
add_two / Jax / tpu / Primal 0.000001470525 s 0.00000149 s 0.99
add_two / HLOOpt / tpu / Primal 0.0000014263000000000002 s 0.00000143465 s 0.99
add_two / PartOpt / tpu / Primal 0.000001472975 s 0.00000148095 s 0.99
add_two / IPartOpt / tpu / Primal 0.00000142935 s 0.000001428325 s 1.00
add_two / DefOpt / tpu / Primal 0.0000014746999999999998 s 0.00000147705 s 1.00
add_two / IDefOpt / tpu / Primal 0.00000143615 s 0.0000014395 s 1.00
add_two / JaXPipe / tpu / Forward 0.00000182505 s 0.0000018262 s 1.00
add_two / Jax / tpu / Forward 0.000001835675 s 0.0000018269 s 1.00
add_two / HLOOpt / tpu / Forward 0.00000182435 s 0.00000182515 s 1.00
add_two / PartOpt / tpu / Forward 0.00000182535 s 0.00000182465 s 1.00
add_two / IPartOpt / tpu / Forward 0.00000182265 s 0.0000018289 s 1.00
add_two / DefOpt / tpu / Forward 0.000001836575 s 0.00000183425 s 1.00
add_two / IDefOpt / tpu / Forward 0.000001829025 s 0.000001825775 s 1.00
add_two / JaXPipe / tpu / PreRev 0.0000028436 s 0.0000028433750000000005 s 1.00
add_two / JaXPipe / tpu / PostRev 0.0000027582250000000006 s 0.0000027536500000000003 s 1.00
add_two / JaXPipe / tpu / BothRev 0.0000028321 s 0.00000285435 s 0.99
add_two / Jax / tpu / BothRev 0.000002744425 s 0.0000027592 s 0.99
add_two / HLOOpt / tpu / PreRev 0.000002826375 s 0.000002848575 s 0.99
add_two / HLOOpt / tpu / PostRev 0.0000027442 s 0.0000027667000000000005 s 0.99
add_two / HLOOpt / tpu / BothRev 0.0000028357 s 0.00000284425 s 1.00
add_two / PartOpt / tpu / PreRev 0.000002745625 s 0.000002759825 s 0.99
add_two / PartOpt / tpu / PostRev 0.0000028365499999999995 s 0.000002855175 s 0.99
add_two / PartOpt / tpu / BothRev 0.0000027564250000000005 s 0.000002765525 s 1.00
add_two / IPartOpt / tpu / PreRev 0.0000028447 s 0.00000285255 s 1.00
add_two / IPartOpt / tpu / PostRev 0.0000027482 s 0.0000027584999999999995 s 1.00
add_two / IPartOpt / tpu / BothRev 0.00000284355 s 0.000002836975 s 1.00
add_two / DefOpt / tpu / PreRev 0.00000275805 s 0.00000276275 s 1.00
add_two / DefOpt / tpu / PostRev 0.000002838775 s 0.00000285745 s 0.99
add_two / DefOpt / tpu / BothRev 0.000002754675 s 0.000002763975 s 1.00
add_two / IDefOpt / tpu / PreRev 0.000002841975 s 0.0000028483000000000003 s 1.00
add_two / IDefOpt / tpu / PostRev 0.00000275485 s 0.0000027806 s 0.99
add_two / IDefOpt / tpu / BothRev 0.000002843325 s 0.000002859375 s 0.99
add_two / JaXPipe / cpu / Primal 0.000013482 s 0.000006799080001655966 s 1.98
add_two / Jax / cpu / Primal 0.000013434 s 0.000007222659960461897 s 1.86
add_two / HLOOpt / cpu / Primal 0.000013263 s 0.000006997399977990426 s 1.90
add_two / PartOpt / cpu / Primal 0.000013391 s 0.000006764840018149698 s 1.98
add_two / IPartOpt / cpu / Primal 0.000013242 s 0.000007197699969765381 s 1.84
add_two / DefOpt / cpu / Primal 0.000013389 s 0.00000721067998711078 s 1.86
add_two / IDefOpt / cpu / Primal 0.00001346 s 0.000006686059987259796 s 2.01
add_two / JaXPipe / cpu / Forward 0.000018267 s 0.00000990931998785527 s 1.84
add_two / Jax / cpu / Forward 0.000018183 s 0.000010424320007587083 s 1.74
add_two / HLOOpt / cpu / Forward 0.00001791 s 0.000010333659993193578 s 1.73
add_two / PartOpt / cpu / Forward 0.000017862 s 0.000010246299971186093 s 1.74
add_two / IPartOpt / cpu / Forward 0.000018069000000000003 s 0.000010061119974125175 s 1.80
add_two / DefOpt / cpu / Forward 0.000018182 s 0.00001001847997031291 s 1.81
add_two / IDefOpt / cpu / Forward 0.000017936 s 0.000010443999999552032 s 1.72
add_two / JaXPipe / cpu / PreRev 0.000023428 s 0.000014581140003429028 s 1.61
add_two / JaXPipe / cpu / PostRev 0.000023252 s 0.000014217040024959715 s 1.64
add_two / JaXPipe / cpu / BothRev 0.000023258 s 0.000014091619987084412 s 1.65
add_two / Jax / cpu / BothRev 0.000023079 s 0.000014226019975467351 s 1.62
add_two / HLOOpt / cpu / PreRev 0.000023302 s 0.0000144150200412696 s 1.62
add_two / HLOOpt / cpu / PostRev 0.000023359 s 0.000018896840019806403 s 1.24
add_two / HLOOpt / cpu / BothRev 0.000023733 s 0.000013867099996787146 s 1.71
add_two / PartOpt / cpu / PreRev 0.000022553 s 0.000014706779975313113 s 1.53
add_two / PartOpt / cpu / PostRev 0.000023133 s 0.000014247260014599306 s 1.62
add_two / PartOpt / cpu / BothRev 0.000023119 s 0.000014383599991560911 s 1.61
add_two / IPartOpt / cpu / PreRev 0.000022870000000000003 s 0.00001433168001312879 s 1.60
add_two / IPartOpt / cpu / PostRev 0.000023328 s 0.000014152100011415314 s 1.65
add_two / IPartOpt / cpu / BothRev 0.000023289 s 0.000014447140029005822 s 1.61
add_two / DefOpt / cpu / PreRev 0.000023673 s 0.000014267520018620417 s 1.66
add_two / DefOpt / cpu / PostRev 0.000023303 s 0.000014259899999160552 s 1.63
add_two / DefOpt / cpu / BothRev 0.000023094 s 0.000014279120023275029 s 1.62
add_two / IDefOpt / cpu / PreRev 0.00002328 s 0.00001455288002944144 s 1.60
add_two / IDefOpt / cpu / PostRev 0.000023466 s 0.000014553979999618603 s 1.61
add_two / IDefOpt / cpu / BothRev 0.000022966 s 0.000014658500022051158 s 1.57
cache / JaXPipe / cpu / Primal 0.000006150160043034703 s 0.000006307679987003212 s 0.98
cache / Jax / cpu / Primal 0.000006544699954247335 s 0.000006591660003323341 s 0.99
cache / HLOOpt / cpu / Primal 0.000006102720017224783 s 0.000006827339984738501 s 0.89
cache / PartOpt / cpu / Primal 0.000005974119958409574 s 0.000006133799961389741 s 0.97
cache / IPartOpt / cpu / Primal 0.000006583280010090675 s 0.000006438860027628834 s 1.02
cache / DefOpt / cpu / Primal 0.000006485079920821591 s 0.000006423040031222627 s 1.01
cache / IDefOpt / cpu / Primal 0.000006423039940273156 s 0.000005969559988443507 s 1.08
cache / JaXPipe / cpu / Forward 0.000014090160093473968 s 0.00001563903999340255 s 0.90
cache / Jax / cpu / Forward 0.000014270779975049665 s 0.000014838980005151826 s 0.96
cache / HLOOpt / cpu / Forward 0.00001530478006316116 s 0.000015777540002090974 s 0.97
cache / PartOpt / cpu / Forward 0.000013985579989821418 s 0.000015409440002258635 s 0.91
cache / IPartOpt / cpu / Forward 0.000015571700059808792 s 0.000015998959997887142 s 0.97
cache / DefOpt / cpu / Forward 0.0000146488400605449 s 0.000014683400013382198 s 1.00
cache / IDefOpt / cpu / Forward 0.000014265859936131163 s 0.00001531049999357492 s 0.93
cache / JaXPipe / cpu / PreRev 0.00001549148011690704 s 0.000015834379964871915 s 0.98
cache / JaXPipe / cpu / PostRev 0.000021110180096002296 s 0.00002103547999467992 s 1.00
cache / JaXPipe / cpu / BothRev 0.000016611300015938468 s 0.000017436600028304384 s 0.95
cache / Jax / cpu / BothRev 0.000020546319956338265 s 0.00003770258001168258 s 0.54
cache / HLOOpt / cpu / PreRev 0.000016083280061138795 s 0.00001725681996504136 s 0.93
cache / HLOOpt / cpu / PostRev 0.00001744870003676624 s 0.00002184574001148576 s 0.80
cache / HLOOpt / cpu / BothRev 0.000015304620010283543 s 0.00001554833997033711 s 0.98
cache / PartOpt / cpu / PreRev 0.00001546482004414429 s 0.000015944460037644604 s 0.97
cache / PartOpt / cpu / PostRev 0.00002108562001012615 s 0.00002022903999204573 s 1.04
cache / PartOpt / cpu / BothRev 0.000016120199998113094 s 0.000016357699996660814 s 0.99
cache / IPartOpt / cpu / PreRev 0.000016047300032369094 s 0.000015744579995953246 s 1.02
cache / IPartOpt / cpu / PostRev 0.00002111022000462981 s 0.00002687556004275393 s 0.79
cache / IPartOpt / cpu / BothRev 0.000016334519950760297 s 0.000016836560062074567 s 0.97
cache / DefOpt / cpu / PreRev 0.00001600249999683001 s 0.0000159212599737657 s 1.01
cache / DefOpt / cpu / PostRev 0.000015720539904577892 s 0.000016033699976105707 s 0.98
cache / DefOpt / cpu / BothRev 0.000016176519875443773 s 0.00001639453998905083 s 0.99
cache / IDefOpt / cpu / PreRev 0.000016080520017567322 s 0.00001627468000151566 s 0.99
cache / IDefOpt / cpu / PostRev 0.0000162226000247756 s 0.00001735963992359757 s 0.93
cache / IDefOpt / cpu / BothRev 0.000016281820007861826 s 0.000016383999964091344 s 0.99
cache / JaXPipe / cuda / Primal 0.000002303 s 0.000002304 s 1.00
cache / Jax / cuda / Primal 0.000002303 s 0.000002304 s 1.00
cache / HLOOpt / cuda / Primal 0.00000224 s 0.00000224 s 1
cache / PartOpt / cuda / Primal 0.00000224 s 0.00000224 s 1
cache / IPartOpt / cuda / Primal 0.000002303 s 0.000002335 s 0.99
cache / DefOpt / cuda / Primal 0.00000224 s 0.000002273 s 0.99
cache / IDefOpt / cuda / Primal 0.000002208 s 0.000002272 s 0.97
cache / JaXPipe / cuda / Forward 0.000002336 s 0.000002335 s 1.00
cache / Jax / cuda / Forward 0.000002304 s 0.000002335 s 0.99
cache / HLOOpt / cuda / Forward 0.000002335 s 0.000002335 s 1
cache / PartOpt / cuda / Forward 0.000002335 s 0.000002335 s 1
cache / IPartOpt / cuda / Forward 0.000002335 s 0.000002335 s 1
cache / DefOpt / cuda / Forward 0.000002304 s 0.000002272 s 1.01
cache / IDefOpt / cuda / Forward 0.000002336 s 0.000002336 s 1
cache / JaXPipe / cuda / PreRev 0.000010752 s 0.00001088 s 0.99
cache / JaXPipe / cuda / PostRev 0.000010433 s 0.000011008 s 0.95
cache / JaXPipe / cuda / BothRev 0.000010657 s 0.000010591 s 1.01
cache / Jax / cuda / BothRev 0.000010336 s 0.000010975 s 0.94
cache / HLOOpt / cuda / PreRev 0.000013439 s 0.000013408 s 1.00
cache / HLOOpt / cuda / PostRev 0.000013408 s 0.000013409000000000002 s 1.00
cache / HLOOpt / cuda / BothRev 0.000013408 s 0.000013408 s 1
cache / PartOpt / cuda / PreRev 0.000010848 s 0.00001088 s 1.00
cache / PartOpt / cuda / PostRev 0.000010784 s 0.00001104 s 0.98
cache / PartOpt / cuda / BothRev 0.000010752 s 0.000011007 s 0.98
cache / IPartOpt / cuda / PreRev 0.000010687 s 0.000010688 s 1.00
cache / IPartOpt / cuda / PostRev 0.000010688 s 0.000010944 s 0.98
cache / IPartOpt / cuda / BothRev 0.00001072 s 0.000010912 s 0.98
cache / DefOpt / cuda / PreRev 0.000010784 s 0.000010816 s 1.00
cache / DefOpt / cuda / PostRev 0.00001088 s 0.000010784 s 1.01
cache / DefOpt / cuda / BothRev 0.000010624 s 0.000010785 s 0.99
cache / IDefOpt / cuda / PreRev 0.00001072 s 0.000010912 s 0.98
cache / IDefOpt / cuda / PostRev 0.00001056 s 0.00001088 s 0.97
cache / IDefOpt / cuda / BothRev 0.000010848 s 0.000010688 s 1.01
cache / JaXPipe / tpu / Primal 0.000002472625 s 0.00000245085 s 1.01
cache / Jax / tpu / Primal 0.0000024525250000000003 s 0.0000024729 s 0.99
cache / HLOOpt / tpu / Primal 0.00000244305 s 0.0000024556500000000004 s 0.99
cache / PartOpt / tpu / Primal 0.00000246615 s 0.0000024559 s 1.00
cache / IPartOpt / tpu / Primal 0.000002465475 s 0.0000024524 s 1.01
cache / DefOpt / tpu / Primal 0.0000024421 s 0.0000024672000000000003 s 0.99
cache / IDefOpt / tpu / Primal 0.00000246 s 0.000002476925 s 0.99
cache / JaXPipe / tpu / Forward 0.00000354035 s 0.0000035351500000000004 s 1.00
cache / Jax / tpu / Forward 0.000003536175 s 0.0000035542 s 0.99
cache / HLOOpt / tpu / Forward 0.0000035606 s 0.00000355125 s 1.00
cache / PartOpt / tpu / Forward 0.0000035303 s 0.000003535775 s 1.00
cache / IPartOpt / tpu / Forward 0.000003545475 s 0.0000035647250000000004 s 0.99
cache / DefOpt / tpu / Forward 0.00000353165 s 0.000003530025 s 1.00
cache / IDefOpt / tpu / Forward 0.000003547425 s 0.00000355895 s 1.00
cache / JaXPipe / tpu / PreRev 0.000004966575 s 0.000004948925 s 1.00
cache / JaXPipe / tpu / PostRev 0.0000049633 s 0.000004943575 s 1.00
cache / JaXPipe / tpu / BothRev 0.000004964025 s 0.0000049926 s 0.99
cache / Jax / tpu / BothRev 0.00000495615 s 0.0000049899 s 0.99
cache / HLOOpt / tpu / PreRev 0.0000039346500000000005 s 0.000003952875 s 1.00
cache / HLOOpt / tpu / PostRev 0.000004118250000000001 s 0.000004132775 s 1.00
cache / HLOOpt / tpu / BothRev 0.00000395425 s 0.000003959050000000001 s 1.00
cache / PartOpt / tpu / PreRev 0.00000497825 s 0.0000049826 s 1.00
cache / PartOpt / tpu / PostRev 0.0000049778 s 0.00000497455 s 1.00
cache / PartOpt / tpu / BothRev 0.000004962849999999999 s 0.00000497915 s 1.00
cache / IPartOpt / tpu / PreRev 0.0000049872 s 0.0000049774 s 1.00
cache / IPartOpt / tpu / PostRev 0.0000049829 s 0.0000049795500000000005 s 1.00
cache / IPartOpt / tpu / BothRev 0.000004956975 s 0.000004968025 s 1.00
cache / DefOpt / tpu / PreRev 0.000004974250000000001 s 0.0000049699000000000005 s 1.00
cache / DefOpt / tpu / PostRev 0.000004974025 s 0.0000049801 s 1.00
cache / DefOpt / tpu / BothRev 0.000004969700000000001 s 0.00000495775 s 1.00
cache / IDefOpt / tpu / PreRev 0.000004971900000000001 s 0.000004973275 s 1.00
cache / IDefOpt / tpu / PostRev 0.00000499025 s 0.000004974475 s 1.00
cache / IDefOpt / tpu / BothRev 0.000004977075 s 0.000004971775 s 1.00
cache / JaXPipe / cpu / Primal 0.0000129 s 0.000006307679987003212 s 2.05
cache / Jax / cpu / Primal 0.000012887 s 0.000006591660003323341 s 1.96
cache / HLOOpt / cpu / Primal 0.000012787 s 0.000006827339984738501 s 1.87
cache / PartOpt / cpu / Primal 0.000012948 s 0.000006133799961389741 s 2.11
cache / IPartOpt / cpu / Primal 0.00001288 s 0.000006438860027628834 s 2.00
cache / DefOpt / cpu / Primal 0.000012755 s 0.000006423040031222627 s 1.99
cache / IDefOpt / cpu / Primal 0.00001275 s 0.000005969559988443507 s 2.14
cache / JaXPipe / cpu / Forward 0.000017087 s 0.00001563903999340255 s 1.09
cache / Jax / cpu / Forward 0.000016486 s 0.000014838980005151826 s 1.11
cache / HLOOpt / cpu / Forward 0.000017133 s 0.000015777540002090974 s 1.09
cache / PartOpt / cpu / Forward 0.000017230999999999998 s 0.000015409440002258635 s 1.12
cache / IPartOpt / cpu / Forward 0.000017267000000000003 s 0.000015998959997887142 s 1.08
cache / DefOpt / cpu / Forward 0.000017173 s 0.000014683400013382198 s 1.17
cache / IDefOpt / cpu / Forward 0.000017203 s 0.00001531049999357492 s 1.12
cache / JaXPipe / cpu / PreRev 0.000017915 s 0.000015834379964871915 s 1.13
cache / JaXPipe / cpu / PostRev 0.000019885 s 0.00002103547999467992 s 0.95
cache / JaXPipe / cpu / BothRev 0.000018072 s 0.000017436600028304384 s 1.04
cache / Jax / cpu / BothRev 0.000019863 s 0.00003770258001168258 s 0.53
cache / HLOOpt / cpu / PreRev 0.000017454999999999998 s 0.00001725681996504136 s 1.01
cache / HLOOpt / cpu / PostRev 0.00001728 s 0.00002184574001148576 s 0.79
cache / HLOOpt / cpu / BothRev 0.000017601 s 0.00001554833997033711 s 1.13
cache / PartOpt / cpu / PreRev 0.000017513 s 0.000015944460037644604 s 1.10
cache / PartOpt / cpu / PostRev 0.000019812 s 0.00002022903999204573 s 0.98
cache / PartOpt / cpu / BothRev 0.000017632 s 0.000016357699996660814 s 1.08
cache / IPartOpt / cpu / PreRev 0.000017289 s 0.000015744579995953246 s 1.10
cache / IPartOpt / cpu / PostRev 0.000020351 s 0.00002687556004275393 s 0.76
cache / IPartOpt / cpu / BothRev 0.000017666 s 0.000016836560062074567 s 1.05
cache / DefOpt / cpu / PreRev 0.000018154 s 0.0000159212599737657 s 1.14
cache / DefOpt / cpu / PostRev 0.000017626 s 0.000016033699976105707 s 1.10
cache / DefOpt / cpu / BothRev 0.000024153 s 0.00001639453998905083 s 1.47
cache / IDefOpt / cpu / PreRev 0.000023893 s 0.00001627468000151566 s 1.47
cache / IDefOpt / cpu / PostRev 0.000017474 s 0.00001735963992359757 s 1.01
cache / IDefOpt / cpu / BothRev 0.000017565999999999997 s 0.000016383999964091344 s 1.07
Concat / JaXPipe / cpu / Primal 0.000006589839922526153 s 0.000007008039992797421 s 0.94
Concat / Jax / cpu / Primal 0.000006683060018985998 s 0.000006910560005053412 s 0.97
Concat / HLOOpt / cpu / Primal 0.000006405339991033543 s 0.0000066355800026940416 s 0.97
Concat / PartOpt / cpu / Primal 0.0000062580599478678775 s 0.000006315100017673104 s 0.99
Concat / IPartOpt / cpu / Primal 0.000006559299945365637 s 0.00000672385998768732 s 0.98
Concat / DefOpt / cpu / Primal 0.00000638044000879745 s 0.000006603359997825464 s 0.97
Concat / IDefOpt / cpu / Primal 0.000006527579898829572 s 0.000006776739974156954 s 0.96
Concat / JaXPipe / cpu / Forward 0.00000971251996816136 s 0.00000970459999734885 s 1.00
Concat / Jax / cpu / Forward 0.000010436359989398623 s 0.000009664840008554166 s 1.08
Concat / HLOOpt / cpu / Forward 0.000009953580047294965 s 0.000009945999991032295 s 1.00
Concat / PartOpt / cpu / Forward 0.000010409259975858732 s 0.000010298160013917368 s 1.01
Concat / IPartOpt / cpu / Forward 0.000009940380059560994 s 0.000010229179988527903 s 0.97
Concat / DefOpt / cpu / Forward 0.000009646879916545004 s 0.000009792959981496096 s 0.99
Concat / IDefOpt / cpu / Forward 0.000009719880072225353 s 0.000009998619980251533 s 0.97
Concat / JaXPipe / cpu / PreRev 0.000011005500000464963 s 0.00001109117998566944 s 0.99
Concat / JaXPipe / cpu / PostRev 0.000011273600066488144 s 0.000011185500006831715 s 1.01
Concat / JaXPipe / cpu / BothRev 0.000010921079920080956 s 0.000011662339984468415 s 0.94
Concat / Jax / cpu / BothRev 0.0000114727800610126 s 0.000011574699983611936 s 0.99
Concat / HLOOpt / cpu / PreRev 0.000011457179934950546 s 0.000011854180002046633 s 0.97
Concat / HLOOpt / cpu / PostRev 0.000013182400070945732 s 0.000013122439995640888 s 1.00
Concat / HLOOpt / cpu / BothRev 0.000011018000041076448 s 0.000011466600044514053 s 0.96
Concat / PartOpt / cpu / PreRev 0.000011592080063564936 s 0.000011967220025326242 s 0.97
Concat / PartOpt / cpu / PostRev 0.000011801599903265016 s 0.00001147991998550424 s 1.03
Concat / PartOpt / cpu / BothRev 0.000011624300041148672 s 0.000012181980009700055 s 0.95
Concat / IPartOpt / cpu / PreRev 0.000011367219976818888 s 0.000012034939991281136 s 0.94
Concat / IPartOpt / cpu / PostRev 0.000011359880008967592 s 0.000011964140003328794 s 0.95
Concat / IPartOpt / cpu / BothRev 0.000010563620053289924 s 0.000011217660012334816 s 0.94
Concat / DefOpt / cpu / PreRev 0.00001132126004449674 s 0.000011659739957394775 s 0.97
Concat / DefOpt / cpu / PostRev 0.00001113225995140965 s 0.000011703559994202806 s 0.95
Concat / DefOpt / cpu / BothRev 0.000010740559991972986 s 0.0000116978799997014 s 0.92
Concat / IDefOpt / cpu / PreRev 0.000011596679960348411 s 0.00001161827999567322 s 1.00
Concat / IDefOpt / cpu / PostRev 0.000010963300010189414 s 0.000011777499994423124 s 0.93
Concat / IDefOpt / cpu / BothRev 0.000011301919985271524 s 0.000011760860015783692 s 0.96
Concat / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
Concat / Jax / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
Concat / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
Concat / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
Concat / IPartOpt / cuda / Primal 0.000001919 s 0.0000019200000000000003 s 1.00
Concat / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
Concat / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.0000019200000000000003 s 1
Concat / JaXPipe / cuda / Forward 0.000009856 s 0.000009824 s 1.00
Concat / Jax / cuda / Forward 0.00000992 s 0.000009568 s 1.04
Concat / HLOOpt / cuda / Forward 0.000009824 s 0.000010144 s 0.97
Concat / PartOpt / cuda / Forward 0.000009824 s 0.000009856 s 1.00
Concat / IPartOpt / cuda / Forward 0.000010209 s 0.000010368 s 0.98
Concat / DefOpt / cuda / Forward 0.00001008 s 0.00000992 s 1.02
Concat / IDefOpt / cuda / Forward 0.000010112 s 0.000009888 s 1.02
Concat / JaXPipe / cuda / PreRev 0.000016224 s 0.000016768000000000003 s 0.97
Concat / JaXPipe / cuda / PostRev 0.00001616 s 0.000016831 s 0.96
Concat / JaXPipe / cuda / BothRev 0.000016736 s 0.00001632 s 1.03
Concat / Jax / cuda / BothRev 0.000016416 s 0.000016737 s 0.98
Concat / HLOOpt / cuda / PreRev 0.000016448000000000002 s 0.000016608 s 0.99
Concat / HLOOpt / cuda / PostRev 0.000015872 s 0.00001632 s 0.97
Concat / HLOOpt / cuda / BothRev 0.000016032 s 0.00001616 s 0.99
Concat / PartOpt / cuda / PreRev 0.000016576000000000002 s 0.00001712 s 0.97
Concat / PartOpt / cuda / PostRev 0.000015968 s 0.000016670999999999997 s 0.96
Concat / PartOpt / cuda / BothRev 0.000016448000000000002 s 0.00001696 s 0.97
Concat / IPartOpt / cuda / PreRev 0.00001616 s 0.000017152 s 0.94
Concat / IPartOpt / cuda / PostRev 0.000015745 s 0.000016383999999999998 s 0.96
Concat / IPartOpt / cuda / BothRev 0.000016448000000000002 s 0.000016736 s 0.98
Concat / DefOpt / cuda / PreRev 0.000016544 s 0.000016416 s 1.01
Concat / DefOpt / cuda / PostRev 0.000016096 s 0.000016927999999999998 s 0.95
Concat / DefOpt / cuda / BothRev 0.000016513 s 0.000016927999999999998 s 0.98
Concat / IDefOpt / cuda / PreRev 0.00001568 s 0.000024671 s 0.64
Concat / IDefOpt / cuda / PostRev 0.000016255999999999998 s 0.000016703 s 0.97
Concat / IDefOpt / cuda / BothRev 0.000016288 s 0.00001648 s 0.99
Concat / JaXPipe / tpu / Primal 0.0000015396 s 0.000001524175 s 1.01
Concat / Jax / tpu / Primal 0.00000153195 s 0.0000015307999999999998 s 1.00
Concat / HLOOpt / tpu / Primal 0.00000153885 s 0.000001524975 s 1.01
Concat / PartOpt / tpu / Primal 0.00000154295 s 0.00000153215 s 1.01
Concat / IPartOpt / tpu / Primal 0.000001536075 s 0.00000152025 s 1.01
Concat / DefOpt / tpu / Primal 0.0000015363750000000002 s 0.000001536925 s 1.00
Concat / IDefOpt / tpu / Primal 0.000001537925 s 0.000001526325 s 1.01
Concat / JaXPipe / tpu / Forward 0.000001579875 s 0.0000015745 s 1.00
Concat / Jax / tpu / Forward 0.0000015484 s 0.0000015418 s 1.00
Concat / HLOOpt / tpu / Forward 0.0000015685 s 0.000001583675 s 0.99
Concat / PartOpt / tpu / Forward 0.0000015487 s 0.0000015417 s 1.00
Concat / IPartOpt / tpu / Forward 0.000001566825 s 0.00000158935 s 0.99
Concat / DefOpt / tpu / Forward 0.0000015581 s 0.000001551225 s 1.00
Concat / IDefOpt / tpu / Forward 0.00000156935 s 0.0000015831500000000002 s 0.99
Concat / JaXPipe / tpu / PreRev 0.000002013025 s 0.0000019998 s 1.01
Concat / JaXPipe / tpu / PostRev 0.000002086925 s 0.0000020947 s 1.00
Concat / JaXPipe / tpu / BothRev 0.0000020035 s 0.000002003425 s 1.00
Concat / Jax / tpu / BothRev 0.00000207515 s 0.000002079375 s 1.00
Concat / HLOOpt / tpu / PreRev 0.0000020047 s 0.0000019996249999999995 s 1.00
Concat / HLOOpt / tpu / PostRev 0.00000208025 s 0.000002077325 s 1.00
Concat / HLOOpt / tpu / BothRev 0.000002013525 s 0.0000019949 s 1.01
Concat / PartOpt / tpu / PreRev 0.0000020814750000000003 s 0.00000207675 s 1.00
Concat / PartOpt / tpu / PostRev 0.00000200565 s 0.000001998025 s 1.00
Concat / PartOpt / tpu / BothRev 0.0000020795 s 0.0000020791 s 1.00
Concat / IPartOpt / tpu / PreRev 0.0000020108500000000004 s 0.000001992975 s 1.01
Concat / IPartOpt / tpu / PostRev 0.000002068875 s 0.000002076225 s 1.00
Concat / IPartOpt / tpu / BothRev 0.000002011475 s 0.0000019937 s 1.01
Concat / DefOpt / tpu / PreRev 0.0000020785 s 0.000002078075 s 1.00
Concat / DefOpt / tpu / PostRev 0.00000200595 s 0.00000199665 s 1.00
Concat / DefOpt / tpu / BothRev 0.000002071975 s 0.00000208455 s 0.99
Concat / IDefOpt / tpu / PreRev 0.0000020079250000000003 s 0.0000020115 s 1.00
Concat / IDefOpt / tpu / PostRev 0.0000020769 s 0.00000208505 s 1.00
Concat / IDefOpt / tpu / BothRev 0.00000201475 s 0.0000020053750000000003 s 1.00
Concat / JaXPipe / cpu / Primal 0.000012592 s 0.000007008039992797421 s 1.80
Concat / Jax / cpu / Primal 0.000013001 s 0.000006910560005053412 s 1.88
Concat / HLOOpt / cpu / Primal 0.000012458 s 0.0000066355800026940416 s 1.88
Concat / PartOpt / cpu / Primal 0.000012672 s 0.000006315100017673104 s 2.01
Concat / IPartOpt / cpu / Primal 0.000012605 s 0.00000672385998768732 s 1.87
Concat / DefOpt / cpu / Primal 0.000012713 s 0.000006603359997825464 s 1.93
Concat / IDefOpt / cpu / Primal 0.000012732 s 0.000006776739974156954 s 1.88
Concat / JaXPipe / cpu / Forward 0.000017827 s 0.00000970459999734885 s 1.84
Concat / Jax / cpu / Forward 0.000017846999999999997 s 0.000009664840008554166 s 1.85
Concat / HLOOpt / cpu / Forward 0.000017675 s 0.000009945999991032295 s 1.78
Concat / PartOpt / cpu / Forward 0.000017344 s 0.000010298160013917368 s 1.68
Concat / IPartOpt / cpu / Forward 0.000017576000000000002 s 0.000010229179988527903 s 1.72
Concat / DefOpt / cpu / Forward 0.000017525 s 0.000009792959981496096 s 1.79
Concat / IDefOpt / cpu / Forward 0.00001748 s 0.000009998619980251533 s 1.75
Concat / JaXPipe / cpu / PreRev 0.000020139 s 0.00001109117998566944 s 1.82
Concat / JaXPipe / cpu / PostRev 0.000019998 s 0.000011185500006831715 s 1.79
Concat / JaXPipe / cpu / BothRev 0.000019799 s 0.000011662339984468415 s 1.70
Concat / Jax / cpu / BothRev 0.000019615 s 0.000011574699983611936 s 1.69
Concat / HLOOpt / cpu / PreRev 0.000019916 s 0.000011854180002046633 s 1.68
Concat / HLOOpt / cpu / PostRev 0.000019619 s 0.000013122439995640888 s 1.50
Concat / HLOOpt / cpu / BothRev 0.000019793 s 0.000011466600044514053 s 1.73
Concat / PartOpt / cpu / PreRev 0.00002003 s 0.000011967220025326242 s 1.67
Concat / PartOpt / cpu / PostRev 0.000019789 s 0.00001147991998550424 s 1.72
Concat / PartOpt / cpu / BothRev 0.00002031 s 0.000012181980009700055 s 1.67
Concat / IPartOpt / cpu / PreRev 0.000019966 s 0.000012034939991281136 s 1.66
Concat / IPartOpt / cpu / PostRev 0.000019542 s 0.000011964140003328794 s 1.63
Concat / IPartOpt / cpu / BothRev 0.000019491 s 0.000011217660012334816 s 1.74
Concat / DefOpt / cpu / PreRev 0.000019861 s 0.000011659739957394775 s 1.70
Concat / DefOpt / cpu / PostRev 0.000019944 s 0.000011703559994202806 s 1.70
Concat / DefOpt / cpu / BothRev 0.000019664 s 0.0000116978799997014 s 1.68
Concat / IDefOpt / cpu / PreRev 0.000019751 s 0.00001161827999567322 s 1.70
Concat / IDefOpt / cpu / PostRev 0.000019803 s 0.000011777499994423124 s 1.68
Concat / IDefOpt / cpu / BothRev 0.00001966 s 0.000011760860015783692 s 1.67
const_scatter / JaXPipe / cpu / Primal 0.000006203100019774865 s 0.0000062349000108952165 s 0.99
const_scatter / Jax / cpu / Primal 0.000006285079980443697 s 0.00000655261999781942 s 0.96
const_scatter / HLOOpt / cpu / Primal 0.000007226979996630689 s 0.000007446100034940173 s 0.97
const_scatter / PartOpt / cpu / Primal 0.000006029900050634751 s 0.000006509699978778372 s 0.93
const_scatter / IPartOpt / cpu / Primal 0.000006507419948320603 s 0.000006898999972690945 s 0.94
const_scatter / DefOpt / cpu / Primal 0.000007306160077860113 s 0.0000069046599946887 s 1.06
const_scatter / IDefOpt / cpu / Primal 0.000007105620061338413 s 0.00000668348001454433 s 1.06
const_scatter / JaXPipe / cpu / Forward 0.000010193940088356612 s 0.00001064506002876442 s 0.96
const_scatter / Jax / cpu / Forward 0.000009089799932553432 s 0.00000971563999883074 s 0.94
const_scatter / HLOOpt / cpu / Forward 0.000010419680002087262 s 0.00001086616000975482 s 0.96
const_scatter / PartOpt / cpu / Forward 0.00001048323998475098 s 0.00001097888003641856 s 0.95
const_scatter / IPartOpt / cpu / Forward 0.000011918839973077411 s 0.000011530499987202348 s 1.03
const_scatter / DefOpt / cpu / Forward 0.000014085760012676474 s 0.00001086903999748756 s 1.30
const_scatter / IDefOpt / cpu / Forward 0.00001045655999405426 s 0.00001085090003471123 s 0.96
const_scatter / JaXPipe / cpu / PreRev 0.000286625040062 s 0.0003115300599802 s 0.92
const_scatter / JaXPipe / cpu / PostRev 0.0002802973400321 s 0.0002813080199848 s 1.00
const_scatter / JaXPipe / cpu / BothRev 0.0002824806399803 s 0.0002825666399712 s 1.00
const_scatter / Jax / cpu / BothRev 0.0002833454199208 s 0.0002801638399796 s 1.01
const_scatter / HLOOpt / cpu / PreRev 0.0002838545999657 s 0.0002834601199992 s 1.00
const_scatter / HLOOpt / cpu / PostRev 0.0002848959799848 s 0.0002848288000222 s 1.00
const_scatter / HLOOpt / cpu / BothRev 0.0002861111599668 s 0.0002825747000315 s 1.01
const_scatter / PartOpt / cpu / PreRev 0.0002833728600126 s 0.0002811166400533 s 1.01
const_scatter / PartOpt / cpu / PostRev 0.0002815997799916 s 0.0002791229000013 s 1.01
const_scatter / PartOpt / cpu / BothRev 0.0002819982399887 s 0.000281865579991 s 1.00
const_scatter / IPartOpt / cpu / PreRev 0.0002823400800662 s 0.0002801982199798 s 1.01
const_scatter / IPartOpt / cpu / PostRev 0.0002923364600246 s 0.0002806498799873 s 1.04
const_scatter / IPartOpt / cpu / BothRev 0.0002822433399342 s 0.0002829532800205 s 1.00
const_scatter / DefOpt / cpu / PreRev 0.000283498419958 s 0.0002809657799934 s 1.01
const_scatter / DefOpt / cpu / PostRev 0.000281489000099 s 0.0002825006799594 s 1.00
const_scatter / DefOpt / cpu / BothRev 0.0002820031199553 s 0.0002812875200197 s 1.00
const_scatter / IDefOpt / cpu / PreRev 0.0002821747400048 s 0.0002909929399902 s 0.97
const_scatter / IDefOpt / cpu / PostRev 0.0002828565401068 s 0.0002834680999694 s 1.00
const_scatter / IDefOpt / cpu / BothRev 0.0002817712399701 s 0.0002822982999714 s 1.00
const_scatter / JaXPipe / cuda / Primal 0.000001887 s 0.000001887 s 1
const_scatter / Jax / cuda / Primal 0.000001887 s 0.000001887 s 1
const_scatter / HLOOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
const_scatter / PartOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
const_scatter / IPartOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
const_scatter / DefOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
const_scatter / IDefOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
const_scatter / JaXPipe / cuda / Forward 0.000009632 s 0.000010113 s 0.95
const_scatter / Jax / cuda / Forward 0.000009824 s 0.000009824 s 1
const_scatter / HLOOpt / cuda / Forward 0.000009824 s 0.00000992 s 0.99
const_scatter / PartOpt / cuda / Forward 0.000009696 s 0.000009984 s 0.97
const_scatter / IPartOpt / cuda / Forward 0.000010144 s 0.00000992 s 1.02
const_scatter / DefOpt / cuda / Forward 0.000009792 s 0.000009856 s 0.99
const_scatter / IDefOpt / cuda / Forward 0.000009728 s 0.000009504 s 1.02
const_scatter / JaXPipe / cuda / PreRev 0.000016704 s 0.000016576000000000002 s 1.01
const_scatter / JaXPipe / cuda / PostRev 0.000015808 s 0.00001584 s 1.00
const_scatter / JaXPipe / cuda / BothRev 0.000016352 s 0.000016705 s 0.98
const_scatter / Jax / cuda / BothRev 0.000015935999999999998 s 0.000016864 s 0.94
const_scatter / HLOOpt / cuda / PreRev 0.000016352 s 0.000016737 s 0.98
const_scatter / HLOOpt / cuda / PostRev 0.000015904000000000002 s 0.00001712 s 0.93
const_scatter / HLOOpt / cuda / BothRev 0.000016063999999999997 s 0.000016385 s 0.98
const_scatter / PartOpt / cuda / PreRev 0.000015872 s 0.000016832 s 0.94
const_scatter / PartOpt / cuda / PostRev 0.000016065 s 0.00001568 s 1.02
const_scatter / PartOpt / cuda / BothRev 0.000016192 s 0.000015937 s 1.02
const_scatter / IPartOpt / cuda / PreRev 0.000016289000000000003 s 0.000016768999999999998 s 0.97
const_scatter / IPartOpt / cuda / PostRev 0.000016448000000000002 s 0.000017984 s 0.91
const_scatter / IPartOpt / cuda / BothRev 0.000015968 s 0.000016096 s 0.99
const_scatter / DefOpt / cuda / PreRev 0.000016255999999999998 s 0.000016736 s 0.97
const_scatter / DefOpt / cuda / PostRev 0.00001808 s 0.000016576000000000002 s 1.09
const_scatter / DefOpt / cuda / BothRev 0.000016095 s 0.000017312 s 0.93
const_scatter / IDefOpt / cuda / PreRev 0.000016 s 0.000016705 s 0.96
const_scatter / IDefOpt / cuda / PostRev 0.000015935000000000002 s 0.000016608 s 0.96
const_scatter / IDefOpt / cuda / BothRev 0.000016224 s 0.000016257 s 1.00
const_scatter / JaXPipe / tpu / Primal 0.0000037984 s 0.000003827175 s 0.99
const_scatter / Jax / tpu / Primal 0.000003824225 s 0.000003828075000000001 s 1.00
const_scatter / HLOOpt / tpu / Primal 0.000003797825 s 0.00000381715 s 0.99
const_scatter / PartOpt / tpu / Primal 0.000003812975 s 0.00000383065 s 1.00
const_scatter / IPartOpt / tpu / Primal 0.000003805975 s 0.0000038145 s 1.00
const_scatter / DefOpt / tpu / Primal 0.0000038034 s 0.000003827675 s 0.99
const_scatter / IDefOpt / tpu / Primal 0.000003806800000000001 s 0.00000379325 s 1.00
const_scatter / JaXPipe / tpu / Forward 0.000006472400000000001 s 0.000006433525 s 1.01
const_scatter / Jax / tpu / Forward 0.000006484975 s 0.000006517650000000001 s 0.99
const_scatter / HLOOpt / tpu / Forward 0.000006448575 s 0.0000064666 s 1.00
const_scatter / PartOpt / tpu / Forward 0.000006510725 s 0.0000065188 s 1.00
const_scatter / IPartOpt / tpu / Forward 0.000006464299999999999 s 0.000006463575 s 1.00
const_scatter / DefOpt / tpu / Forward 0.000006501400000000001 s 0.0000065058000000000005 s 1.00
const_scatter / IDefOpt / tpu / Forward 0.000006469325000000001 s 0.0000064621 s 1.00
const_scatter / JaXPipe / tpu / PreRev 0.00000663655 s 0.000006682225000000001 s 0.99
const_scatter / JaXPipe / tpu / PostRev 0.00000660665 s 0.0000066495 s 0.99
const_scatter / JaXPipe / tpu / BothRev 0.0000066152 s 0.000006660075 s 0.99
const_scatter / Jax / tpu / BothRev 0.0000066198000000000006 s 0.000006669049999999999 s 0.99
const_scatter / HLOOpt / tpu / PreRev 0.000006619224999999999 s 0.0000066533750000000005 s 0.99
const_scatter / HLOOpt / tpu / PostRev 0.0000066072 s 0.000006653775 s 0.99
const_scatter / HLOOpt / tpu / BothRev 0.00000662815 s 0.000006698075 s 0.99
const_scatter / PartOpt / tpu / PreRev 0.000006639549999999999 s 0.000006674774999999999 s 0.99
const_scatter / PartOpt / tpu / PostRev 0.0000065976 s 0.000006659625 s 0.99
const_scatter / PartOpt / tpu / BothRev 0.000006613750000000001 s 0.0000066755 s 0.99
const_scatter / IPartOpt / tpu / PreRev 0.0000066031750000000005 s 0.000006695950000000001 s 0.99
const_scatter / IPartOpt / tpu / PostRev 0.0000066308 s 0.000006663375 s 1.00
const_scatter / IPartOpt / tpu / BothRev 0.000006633325 s 0.000006657750000000001 s 1.00
const_scatter / DefOpt / tpu / PreRev 0.000006625525 s 0.00000665115 s 1.00
const_scatter / DefOpt / tpu / PostRev 0.000006622525000000001 s 0.0000066825 s 0.99
const_scatter / DefOpt / tpu / BothRev 0.00000662365 s 0.000006679575 s 0.99
const_scatter / IDefOpt / tpu / PreRev 0.000006613375 s 0.000006672949999999999 s 0.99
const_scatter / IDefOpt / tpu / PostRev 0.000006616575 s 0.0000066757750000000005 s 0.99
const_scatter / IDefOpt / tpu / BothRev 0.000006605775 s 0.00000667055 s 0.99
const_scatter / JaXPipe / cpu / Primal 0.000013004 s 0.0000062349000108952165 s 2.09
const_scatter / Jax / cpu / Primal 0.000012548 s 0.00000655261999781942 s 1.91
const_scatter / HLOOpt / cpu / Primal 0.000013122 s 0.000007446100034940173 s 1.76
const_scatter / PartOpt / cpu / Primal 0.00001275 s 0.000006509699978778372 s 1.96
const_scatter / IPartOpt / cpu / Primal 0.000012621 s 0.000006898999972690945 s 1.83
const_scatter / DefOpt / cpu / Primal 0.000013543 s 0.0000069046599946887 s 1.96
const_scatter / IDefOpt / cpu / Primal 0.000013018 s 0.00000668348001454433 s 1.95
const_scatter / JaXPipe / cpu / Forward 0.00001814 s 0.00001064506002876442 s 1.70
const_scatter / Jax / cpu / Forward 0.000016709000000000002 s 0.00000971563999883074 s 1.72
const_scatter / HLOOpt / cpu / Forward 0.00001835 s 0.00001086616000975482 s 1.69
const_scatter / PartOpt / cpu / Forward 0.000018168 s 0.00001097888003641856 s 1.65
const_scatter / IPartOpt / cpu / Forward 0.00001815 s 0.000011530499987202348 s 1.57
const_scatter / DefOpt / cpu / Forward 0.000017983 s 0.00001086903999748756 s 1.65
const_scatter / IDefOpt / cpu / Forward 0.000017746000000000003 s 0.00001085090003471123 s 1.64
const_scatter / JaXPipe / cpu / PreRev 0.000509542 s 0.0003115300599802 s 1.64
const_scatter / JaXPipe / cpu / PostRev 0.000487853 s 0.0002813080199848 s 1.73
const_scatter / JaXPipe / cpu / BothRev 0.000502687 s 0.0002825666399712 s 1.78
const_scatter / Jax / cpu / BothRev 0.000496507 s 0.0002801638399796 s 1.77
const_scatter / HLOOpt / cpu / PreRev 0.000501453 s 0.0002834601199992 s 1.77
const_scatter / HLOOpt / cpu / PostRev 0.00049199 s 0.0002848288000222 s 1.73
const_scatter / HLOOpt / cpu / BothRev 0.000505775 s 0.0002825747000315 s 1.79
const_scatter / PartOpt / cpu / PreRev 0.00052195 s 0.0002811166400533 s 1.86
const_scatter / PartOpt / cpu / PostRev 0.000485297 s 0.0002791229000013 s 1.74
const_scatter / PartOpt / cpu / BothRev 0.000505306 s 0.000281865579991 s 1.79
const_scatter / IPartOpt / cpu / PreRev 0.000495186 s 0.0002801982199798 s 1.77
const_scatter / IPartOpt / cpu / PostRev 0.000494036 s 0.0002806498799873 s 1.76
const_scatter / IPartOpt / cpu / BothRev 0.000498285 s 0.0002829532800205 s 1.76
const_scatter / DefOpt / cpu / PreRev 0.000509498 s 0.0002809657799934 s 1.81
const_scatter / DefOpt / cpu / PostRev 0.000508788 s 0.0002825006799594 s 1.80
const_scatter / DefOpt / cpu / BothRev 0.000502391 s 0.0002812875200197 s 1.79
const_scatter / IDefOpt / cpu / PreRev 0.000490828 s 0.0002909929399902 s 1.69
const_scatter / IDefOpt / cpu / PostRev 0.0005085819999999 s 0.0002834680999694 s 1.79
const_scatter / IDefOpt / cpu / BothRev 0.000502019 s 0.0002822982999714 s 1.78
GenDot / JaXPipe / cpu / Primal 0.000007022140034678159 s 0.000007543960009570583 s 0.93
GenDot / Jax / cpu / Primal 0.000006917459922988201 s 0.00000684336000631447 s 1.01
GenDot / HLOOpt / cpu / Primal 0.000007230179980979301 s 0.000007792519991198787 s 0.93
GenDot / PartOpt / cpu / Primal 0.000006859999939479167 s 0.000006615960010094568 s 1.04
GenDot / IPartOpt / cpu / Primal 0.000006710700035910122 s 0.000006698160013911547 s 1.00
GenDot / DefOpt / cpu / Primal 0.0000073895000241464 s 0.000007053299959807191 s 1.05
GenDot / IDefOpt / cpu / Primal 0.0000067614199178933634 s 0.00000699161998454656 s 0.97
GenDot / JaXPipe / cpu / Forward 0.000010699939957703464 s 0.000010652839973772645 s 1.00
GenDot / Jax / cpu / Forward 0.00001001787992208847 s 0.000010071080032503232 s 0.99
GenDot / HLOOpt / cpu / Forward 0.000011368659979780204 s 0.000011345020020598896 s 1.00
GenDot / PartOpt / cpu / Forward 0.00001077698001608951 s 0.000010618520036587142 s 1.01
GenDot / IPartOpt / cpu / Forward 0.000010959839946735884 s 0.000011141940003653872 s 0.98
GenDot / DefOpt / cpu / Forward 0.00001024000008328585 s 0.000011015899963240372 s 0.93
GenDot / IDefOpt / cpu / Forward 0.000010169880024477608 s 0.000010992019988407264 s 0.93
GenDot / JaXPipe / cpu / PreRev 0.000011207679999643006 s 0.000011150480022479314 s 1.01
GenDot / JaXPipe / cpu / PostRev 0.000011010859980160603 s 0.0000100475599992933 s 1.10
GenDot / JaXPipe / cpu / BothRev 0.000010801379994518357 s 0.000011710220014720108 s 0.92
GenDot / Jax / cpu / BothRev 0.000009763080070115394 s 0.000010871300009966944 s 0.90
GenDot / HLOOpt / cpu / PreRev 0.000011509540036058752 s 0.00001220950001879828 s 0.94
GenDot / HLOOpt / cpu / PostRev 0.000012955839974893024 s 0.000013349500022741267 s 0.97
GenDot / HLOOpt / cpu / BothRev 0.000010846880068129394 s 0.000011031640015062296 s 0.98
GenDot / PartOpt / cpu / PreRev 0.000011018079985660734 s 0.000010969980039590154 s 1.00
GenDot / PartOpt / cpu / PostRev 0.00001045155999236158 s 0.000010812399959831965 s 0.97
GenDot / PartOpt / cpu / BothRev 0.000011244220040680376 s 0.000011814540011982898 s 0.95
GenDot / IPartOpt / cpu / PreRev 0.000012082619959983277 s 0.000010740520019680844 s 1.12
GenDot / IPartOpt / cpu / PostRev 0.000011159199966641608 s 0.00000975463995928294 s 1.14
GenDot / IPartOpt / cpu / BothRev 0.000011393220047466457 s 0.000011324979986966356 s 1.01
GenDot / DefOpt / cpu / PreRev 0.0000111574400216341 s 0.000011239259993089946 s 0.99
GenDot / DefOpt / cpu / PostRev 0.000010974920041917355 s 0.00001121968001825735 s 0.98
GenDot / DefOpt / cpu / BothRev 0.000011212839945073938 s 0.000011404920005588792 s 0.98
GenDot / IDefOpt / cpu / PreRev 0.000010678300022846089 s 0.0000108012400323787 s 0.99
GenDot / IDefOpt / cpu / PostRev 0.000011301160029688617 s 0.000012132959973314428 s 0.93
GenDot / IDefOpt / cpu / BothRev 0.000010842679930647137 s 0.00001086260001102346 s 1.00
GenDot / JaXPipe / cuda / Primal 0.000002015 s 0.000002015 s 1
GenDot / Jax / cuda / Primal 0.000002016 s 0.000002016 s 1
GenDot / HLOOpt / cuda / Primal 0.000001984 s 0.000001984 s 1
GenDot / PartOpt / cuda / Primal 0.000002015 s 0.000002016 s 1.00
GenDot / IPartOpt / cuda / Primal 0.000002015 s 0.000002016 s 1.00
GenDot / DefOpt / cuda / Primal 0.000002015 s 0.000002016 s 1.00
GenDot / IDefOpt / cuda / Primal 0.000002015 s 0.000002015 s 1
GenDot / JaXPipe / cuda / Forward 0.000010048 s 0.00001024 s 0.98
GenDot / Jax / cuda / Forward 0.00000976 s 0.000009983 s 0.98
GenDot / HLOOpt / cuda / Forward 0.0000112 s 0.000009856 s 1.14
GenDot / PartOpt / cuda / Forward 0.000010367 s 0.000009856 s 1.05
GenDot / IPartOpt / cuda / Forward 0.000010944 s 0.000009824 s 1.11
GenDot / DefOpt / cuda / Forward 0.00001104 s 0.00000992 s 1.11
GenDot / IDefOpt / cuda / Forward 0.000011327 s 0.000009728 s 1.16
GenDot / JaXPipe / cuda / PreRev 0.000009567 s 0.000010112 s 0.95
GenDot / JaXPipe / cuda / PostRev 0.000009569 s 0.000010016 s 0.96
GenDot / JaXPipe / cuda / BothRev 0.000009536 s 0.000010048 s 0.95
GenDot / Jax / cuda / BothRev 0.000009567 s 0.000010144 s 0.94
GenDot / HLOOpt / cuda / PreRev 0.000009728 s 0.000010016 s 0.97
GenDot / HLOOpt / cuda / PostRev 0.000009728 s 0.000010048 s 0.97
GenDot / HLOOpt / cuda / BothRev 0.000009728 s 0.000009664 s 1.01
GenDot / PartOpt / cuda / PreRev 0.000010048 s 0.000010016 s 1.00
GenDot / PartOpt / cuda / PostRev 0.000009856 s 0.00001008 s 0.98
GenDot / PartOpt / cuda / BothRev 0.000009729 s 0.000010592 s 0.92
GenDot / IPartOpt / cuda / PreRev 0.00000976 s 0.000011264 s 0.87
GenDot / IPartOpt / cuda / PostRev 0.000010208 s 0.000011648 s 0.88
GenDot / IPartOpt / cuda / BothRev 0.000009952 s 0.000011328 s 0.88
GenDot / DefOpt / cuda / PreRev 0.000009825 s 0.000011520000000000002 s 0.85
GenDot / DefOpt / cuda / PostRev 0.000009472 s 0.000010527 s 0.90
GenDot / DefOpt / cuda / BothRev 0.000010017 s 0.000011392 s 0.88
GenDot / IDefOpt / cuda / PreRev 0.000009952 s 0.000010048 s 0.99
GenDot / IDefOpt / cuda / PostRev 0.000009951 s 0.000009984 s 1.00
GenDot / IDefOpt / cuda / BothRev 0.000009056 s 0.000009792 s 0.92
GenDot / JaXPipe / tpu / Primal 9.30525e-7 s 9.30525e-7 s 1
GenDot / Jax / tpu / Primal 9.259e-7 s 9.257e-7 s 1.00
GenDot / HLOOpt / tpu / Primal 0.0000015894249999999998 s 0.0000015812249999999998 s 1.01
GenDot / PartOpt / tpu / Primal 9.2495e-7 s 9.25825e-7 s 1.00
GenDot / IPartOpt / tpu / Primal 9.29775e-7 s 9.30225e-7 s 1.00
GenDot / DefOpt / tpu / Primal 0.00000149535 s 0.00000150005 s 1.00
GenDot / IDefOpt / tpu / Primal 0.00000157545 s 0.0000015821 s 1.00
GenDot / JaXPipe / tpu / Forward 0.000003168125 s 0.000003177125 s 1.00
GenDot / Jax / tpu / Forward 0.00000232305 s 0.00000232435 s 1.00
GenDot / HLOOpt / tpu / Forward 0.0000031276250000000004 s 0.00000313075 s 1.00
GenDot / PartOpt / tpu / Forward 0.00000321795 s 0.000003227975 s 1.00
GenDot / IPartOpt / tpu / Forward 0.00000312365 s 0.0000031252 s 1.00
GenDot / DefOpt / tpu / Forward 0.000003215575 s 0.000003226375 s 1.00
GenDot / IDefOpt / tpu / Forward 0.00000312455 s 0.000003131175 s 1.00
GenDot / JaXPipe / tpu / PreRev 0.0000029603000000000003 s 0.0000029877 s 0.99
GenDot / JaXPipe / tpu / PostRev 0.00000240385 s 0.000002404625 s 1.00
GenDot / JaXPipe / tpu / BothRev 0.00000295865 s 0.000002984075 s 0.99
GenDot / Jax / tpu / BothRev 0.000002415875 s 0.0000023980500000000003 s 1.01
GenDot / HLOOpt / tpu / PreRev 0.000002969925 s 0.0000029821500000000003 s 1.00
GenDot / HLOOpt / tpu / PostRev 0.0000029338750000000003 s 0.000002923425 s 1.00
GenDot / HLOOpt / tpu / BothRev 0.00000296385 s 0.0000029864249999999995 s 0.99
GenDot / PartOpt / tpu / PreRev 0.00000294475 s 0.000002925775 s 1.01
GenDot / PartOpt / tpu / PostRev 0.000002408425 s 0.0000023952 s 1.01
GenDot / PartOpt / tpu / BothRev 0.0000029437 s 0.000002923075 s 1.01
GenDot / IPartOpt / tpu / PreRev 0.00000296895 s 0.0000029841 s 0.99
GenDot / IPartOpt / tpu / PostRev 0.000002408275 s 0.000002404425 s 1.00
GenDot / IPartOpt / tpu / BothRev 0.00000295605 s 0.000002975 s 0.99
GenDot / DefOpt / tpu / PreRev 0.0000029367500000000003 s 0.000002927375 s 1.00
GenDot / DefOpt / tpu / PostRev 0.0000029702250000000004 s 0.00000298225 s 1.00
GenDot / DefOpt / tpu / BothRev 0.000002936025 s 0.000002920825 s 1.01
GenDot / IDefOpt / tpu / PreRev 0.0000029672 s 0.00000298085 s 1.00
GenDot / IDefOpt / tpu / PostRev 0.00000293315 s 0.0000029257 s 1.00
GenDot / IDefOpt / tpu / BothRev 0.00000295765 s 0.0000029752250000000004 s 0.99
GenDot / JaXPipe / cpu / Primal 0.000014502 s 0.000007543960009570583 s 1.92
GenDot / Jax / cpu / Primal 0.000015241 s 0.00000684336000631447 s 2.23
GenDot / HLOOpt / cpu / Primal 0.00001398 s 0.000007792519991198787 s 1.79
GenDot / PartOpt / cpu / Primal 0.0000149 s 0.000006615960010094568 s 2.25
GenDot / IPartOpt / cpu / Primal 0.000015307 s 0.000006698160013911547 s 2.29
GenDot / DefOpt / cpu / Primal 0.000014041 s 0.000007053299959807191 s 1.99
GenDot / IDefOpt / cpu / Primal 0.000013993 s 0.00000699161998454656 s 2.00
GenDot / JaXPipe / cpu / Forward 0.000019315 s 0.000010652839973772645 s 1.81
GenDot / Jax / cpu / Forward 0.000020236 s 0.000010071080032503232 s 2.01
GenDot / HLOOpt / cpu / Forward 0.000019054 s 0.000011345020020598896 s 1.68
GenDot / PartOpt / cpu / Forward 0.00001922 s 0.000010618520036587142 s 1.81
GenDot / IPartOpt / cpu / Forward 0.000018853 s 0.000011141940003653872 s 1.69
GenDot / DefOpt / cpu / Forward 0.000019064 s 0.000011015899963240372 s 1.73
GenDot / IDefOpt / cpu / Forward 0.000019534 s 0.000010992019988407264 s 1.78
GenDot / JaXPipe / cpu / PreRev 0.000019266 s 0.000011150480022479314 s 1.73
GenDot / JaXPipe / cpu / PostRev 0.000020206 s 0.0000100475599992933 s 2.01
GenDot / JaXPipe / cpu / BothRev 0.000019091 s 0.000011710220014720108 s 1.63
GenDot / Jax / cpu / BothRev 0.000020421000000000003 s 0.000010871300009966944 s 1.88
GenDot / HLOOpt / cpu / PreRev 0.000018906 s 0.00001220950001879828 s 1.55
GenDot / HLOOpt / cpu / PostRev 0.000019396 s 0.000013349500022741267 s 1.45
GenDot / HLOOpt / cpu / BothRev 0.000019307 s 0.000011031640015062296 s 1.75
GenDot / PartOpt / cpu / PreRev 0.000018876 s 0.000010969980039590154 s 1.72
GenDot / PartOpt / cpu / PostRev 0.000020493 s 0.000010812399959831965 s 1.90
GenDot / PartOpt / cpu / BothRev 0.000019464 s 0.000011814540011982898 s 1.65
GenDot / IPartOpt / cpu / PreRev 0.000019229 s 0.000010740520019680844 s 1.79
GenDot / IPartOpt / cpu / PostRev 0.000020378 s 0.00000975463995928294 s 2.09
GenDot / IPartOpt / cpu / BothRev 0.000019594 s 0.000011324979986966356 s 1.73
GenDot / DefOpt / cpu / PreRev 0.000019072 s 0.000011239259993089946 s 1.70
GenDot / DefOpt / cpu / PostRev 0.000019358 s 0.00001121968001825735 s 1.73
GenDot / DefOpt / cpu / BothRev 0.000019441 s 0.000011404920005588792 s 1.70
GenDot / IDefOpt / cpu / PreRev 0.000018946 s 0.0000108012400323787 s 1.75
GenDot / IDefOpt / cpu / PostRev 0.000019308 s 0.000012132959973314428 s 1.59
GenDot / IDefOpt / cpu / BothRev 0.000019072 s 0.00001086260001102346 s 1.76
hlo_ffi / JaXPipe / cpu / Primal 0.000011587940043682466 s 0.00001027235996843956 s 1.13
hlo_ffi / Jax / cpu / Primal 0.000011073220011894591 s 0.00001028005999614834 s 1.08
hlo_ffi / HLOOpt / cpu / Primal 0.000011319820005155636 s 0.000010306739959560218 s 1.10
hlo_ffi / PartOpt / cpu / Primal 0.000011083279969170687 s 0.000010014420022343984 s 1.11
hlo_ffi / IPartOpt / cpu / Primal 0.00001161695994596812 s 0.000010706400044000476 s 1.09
hlo_ffi / DefOpt / cpu / Primal 0.000010618759934004629 s 0.0000098558800345927 s 1.08
hlo_ffi / IDefOpt / cpu / Primal 0.000010971079955197638 s 0.00000984211998002138 s 1.11
hlo_ffi / JaXPipe / cpu / Forward 0.000016104940041259397 s 0.00001452496000638348 s 1.11
hlo_ffi / Jax / cpu / Forward 0.000015841740059840957 s 0.000014700679985253372 s 1.08
hlo_ffi / HLOOpt / cpu / Forward 0.000016350739988411077 s 0.00001474948001487064 s 1.11
hlo_ffi / PartOpt / cpu / Forward 0.00001672966000114684 s 0.000014515399998344948 s 1.15
hlo_ffi / IPartOpt / cpu / Forward 0.000015834340047149453 s 0.000014186420003170497 s 1.12
hlo_ffi / DefOpt / cpu / Forward 0.00001649398007430136 s 0.000014972280032452544 s 1.10
hlo_ffi / IDefOpt / cpu / Forward 0.000016737200057832523 s 0.000014576040011888836 s 1.15
hlo_ffi / JaXPipe / cpu / PreRev 0.000017899800041050186 s 0.000014949839978726232 s 1.20
hlo_ffi / JaXPipe / cpu / PostRev 0.00001466846002585953 s 0.000014624340046793804 s 1.00
hlo_ffi / JaXPipe / cpu / BothRev 0.00001557003997731954 s 0.000013843159995303725 s 1.12
hlo_ffi / Jax / cpu / BothRev 0.00001645814003495616 s 0.000015019619995655376 s 1.10
hlo_ffi / HLOOpt / cpu / PreRev 0.000016053500021371293 s 0.00001564962006341375 s 1.03
hlo_ffi / HLOOpt / cpu / PostRev 0.000017678860040177824 s 0.000016052320006565424 s 1.10
hlo_ffi / HLOOpt / cpu / BothRev 0.000015761660051794024 s 0.00001432220001333917 s 1.10
hlo_ffi / PartOpt / cpu / PreRev 0.000015957279974827542 s 0.00001576888003910426 s 1.01
hlo_ffi / PartOpt / cpu / PostRev 0.000015466120021301322 s 0.00001419655999598035 s 1.09
hlo_ffi / PartOpt / cpu / BothRev 0.000015454219974344595 s 0.000014534319980157308 s 1.06
hlo_ffi / IPartOpt / cpu / PreRev 0.000016542399989702972 s 0.000015157460038608406 s 1.09
hlo_ffi / IPartOpt / cpu / PostRev 0.000015543639947281917 s 0.000014356760048030991 s 1.08
hlo_ffi / IPartOpt / cpu / BothRev 0.000015074519978952594 s 0.000014438420002989003 s 1.04
hlo_ffi / DefOpt / cpu / PreRev 0.000016861560088727855 s 0.000014575639988834157 s 1.16
hlo_ffi / DefOpt / cpu / PostRev 0.000014954879898141373 s 0.000014225739969333515 s 1.05
hlo_ffi / DefOpt / cpu / BothRev 0.0000153469200085965 s 0.000014324239991765353 s 1.07
hlo_ffi / IDefOpt / cpu / PreRev 0.000016380079941882286 s 0.00001521861997389351 s 1.08
hlo_ffi / IDefOpt / cpu / PostRev 0.000015062300117278937 s 0.000014588500034733444 s 1.03
hlo_ffi / IDefOpt / cpu / BothRev 0.00001510169991888688 s 0.000014252720002332352 s 1.06
hlo_ffi / JaXPipe / cuda / Primal 0.000001984 s 0.000001984 s 1
hlo_ffi / Jax / cuda / Primal 0.000001983 s 0.000001983 s 1
hlo_ffi / HLOOpt / cuda / Primal 0.000001983 s 0.000001984 s 1.00
hlo_ffi / PartOpt / cuda / Primal 0.000001984 s 0.000001984 s 1
hlo_ffi / IPartOpt / cuda / Primal 0.000001983 s 0.000001984 s 1.00
hlo_ffi / DefOpt / cuda / Primal 0.000001984 s 0.000001984 s 1
hlo_ffi / IDefOpt / cuda / Primal 0.000001984 s 0.000001983 s 1.00
hlo_ffi / JaXPipe / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / Jax / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / HLOOpt / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / PartOpt / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / IPartOpt / cuda / Forward 0.000002048 s 0.00000208 s 0.98
hlo_ffi / DefOpt / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / IDefOpt / cuda / Forward 0.00000208 s 0.00000208 s 1
hlo_ffi / JaXPipe / cuda / PreRev 0.000002047 s 0.000002048 s 1.00
hlo_ffi / JaXPipe / cuda / PostRev 0.000002048 s 0.000002048 s 1
hlo_ffi / JaXPipe / cuda / BothRev 0.000002047 s 0.000002047 s 1
hlo_ffi / Jax / cuda / BothRev 0.000002048 s 0.00000208 s 0.98
hlo_ffi / HLOOpt / cuda / PreRev 0.000002047 s 0.000002048 s 1.00
hlo_ffi / HLOOpt / cuda / PostRev 0.000002048 s 0.000002048 s 1
hlo_ffi / HLOOpt / cuda / BothRev 0.000002048 s 0.000002048 s 1
hlo_ffi / PartOpt / cuda / PreRev 0.000002048 s 0.000002048 s 1
hlo_ffi / PartOpt / cuda / PostRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / PartOpt / cuda / BothRev 0.00000208 s 0.000002047 s 1.02
hlo_ffi / IPartOpt / cuda / PreRev 0.000002047 s 0.000002047 s 1
hlo_ffi / IPartOpt / cuda / PostRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / IPartOpt / cuda / BothRev 0.000002048 s 0.000002048 s 1
hlo_ffi / DefOpt / cuda / PreRev 0.000002047 s 0.00000208 s 0.98
hlo_ffi / DefOpt / cuda / PostRev 0.000002047 s 0.000002048 s 1.00
hlo_ffi / DefOpt / cuda / BothRev 0.000002048 s 0.000002048 s 1
hlo_ffi / IDefOpt / cuda / PreRev 0.000002048 s 0.000002047 s 1.00
hlo_ffi / IDefOpt / cuda / PostRev 0.000002048 s 0.000002048 s 1
hlo_ffi / IDefOpt / cuda / BothRev 0.000002048 s 0.000002048 s 1
hlo_ffi / JaXPipe / tpu / Primal 9.31625e-7 s 9.172e-7 s 1.02
hlo_ffi / Jax / tpu / Primal 9.52375e-7 s 9.499e-7 s 1.00
hlo_ffi / HLOOpt / tpu / Primal 9.07e-7 s 8.99325e-7 s 1.01
hlo_ffi / PartOpt / tpu / Primal 9.49425e-7 s 9.55225e-7 s 0.99
hlo_ffi / IPartOpt / tpu / Primal 9.07e-7 s 9.028e-7 s 1.00
hlo_ffi / DefOpt / tpu / Primal 9.541e-7 s 9.5425e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Primal 9.06875e-7 s 8.985500000000001e-7 s 1.01
hlo_ffi / JaXPipe / tpu / Forward 9.493e-7 s 9.49375e-7 s 1.00
hlo_ffi / Jax / tpu / Forward 9.81675e-7 s 9.8175e-7 s 1.00
hlo_ffi / HLOOpt / tpu / Forward 9.7445e-7 s 9.7415e-7 s 1.00
hlo_ffi / PartOpt / tpu / Forward 9.34375e-7 s 9.3425e-7 s 1.00
hlo_ffi / IPartOpt / tpu / Forward 9.73825e-7 s 9.73925e-7 s 1.00
hlo_ffi / DefOpt / tpu / Forward 9.34425e-7 s 9.34025e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Forward 9.74275e-7 s 9.746e-7 s 1.00
hlo_ffi / JaXPipe / tpu / PreRev 9.3785e-7 s 9.323e-7 s 1.01
hlo_ffi / JaXPipe / tpu / PostRev 9.65525e-7 s 9.654e-7 s 1.00
hlo_ffi / JaXPipe / tpu / BothRev 9.6205e-7 s 9.63e-7 s 1.00
hlo_ffi / Jax / tpu / BothRev 9.64975e-7 s 9.653e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PreRev 9.625e-7 s 9.6215e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PostRev 9.648e-7 s 9.65175e-7 s 1.00
hlo_ffi / HLOOpt / tpu / BothRev 9.6265e-7 s 9.62175e-7 s 1.00
hlo_ffi / PartOpt / tpu / PreRev 9.650749999999998e-7 s 9.648e-7 s 1.00
hlo_ffi / PartOpt / tpu / PostRev 9.6245e-7 s 9.62225e-7 s 1.00
hlo_ffi / PartOpt / tpu / BothRev 9.6485e-7 s 9.654e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PreRev 9.62025e-7 s 9.6265e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PostRev 9.6475e-7 s 9.652e-7 s 1.00
hlo_ffi / IPartOpt / tpu / BothRev 9.61825e-7 s 9.619e-7 s 1.00
hlo_ffi / DefOpt / tpu / PreRev 9.651e-7 s 9.6525e-7 s 1.00
hlo_ffi / DefOpt / tpu / PostRev 9.618500000000002e-7 s 9.62225e-7 s 1.00
hlo_ffi / DefOpt / tpu / BothRev 9.6495e-7 s 9.65175e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PreRev 9.61925e-7 s 9.6225e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PostRev 9.6515e-7 s 9.6515e-7 s 1
hlo_ffi / IDefOpt / tpu / BothRev 9.622e-7 s 9.62225e-7 s 1.00
hlo_ffi / JaXPipe / cpu / Primal 0.00001822 s 0.00001027235996843956 s 1.77
hlo_ffi / Jax / cpu / Primal 0.000017801 s 0.00001028005999614834 s 1.73
hlo_ffi / HLOOpt / cpu / Primal 0.000017692 s 0.000010306739959560218 s 1.72
hlo_ffi / PartOpt / cpu / Primal 0.00001749 s 0.000010014420022343984 s 1.75
hlo_ffi / IPartOpt / cpu / Primal 0.000017501 s 0.000010706400044000476 s 1.63
hlo_ffi / DefOpt / cpu / Primal 0.000017760999999999998 s 0.0000098558800345927 s 1.80
hlo_ffi / IDefOpt / cpu / Primal 0.000017701000000000002 s 0.00000984211998002138 s 1.80
hlo_ffi / JaXPipe / cpu / Forward 0.00002467 s 0.00001452496000638348 s 1.70
hlo_ffi / Jax / cpu / Forward 0.000023800000000000003 s 0.000014700679985253372 s 1.62
hlo_ffi / HLOOpt / cpu / Forward 0.000024348 s 0.00001474948001487064 s 1.65
hlo_ffi / PartOpt / cpu / Forward 0.000024049 s 0.000014515399998344948 s 1.66
hlo_ffi / IPartOpt / cpu / Forward 0.000024018000000000003 s 0.000014186420003170497 s 1.69
hlo_ffi / DefOpt / cpu / Forward 0.000024243 s 0.000014972280032452544 s 1.62
hlo_ffi / IDefOpt / cpu / Forward 0.000024052 s 0.000014576040011888836 s 1.65
hlo_ffi / JaXPipe / cpu / PreRev 0.000024408 s 0.000014949839978726232 s 1.63
hlo_ffi / JaXPipe / cpu / PostRev 0.000023554 s 0.000014624340046793804 s 1.61
hlo_ffi / JaXPipe / cpu / BothRev 0.000024105 s 0.000013843159995303725 s 1.74
hlo_ffi / Jax / cpu / BothRev 0.000023849 s 0.000015019619995655376 s 1.59
hlo_ffi / HLOOpt / cpu / PreRev 0.000023873 s 0.00001564962006341375 s 1.53
hlo_ffi / HLOOpt / cpu / PostRev 0.000024034 s 0.000016052320006565424 s 1.50
hlo_ffi / HLOOpt / cpu / BothRev 0.000023942 s 0.00001432220001333917 s 1.67
hlo_ffi / PartOpt / cpu / PreRev 0.000023727 s 0.00001576888003910426 s 1.50
hlo_ffi / PartOpt / cpu / PostRev 0.000024172 s 0.00001419655999598035 s 1.70
hlo_ffi / PartOpt / cpu / BothRev 0.000024442 s 0.000014534319980157308 s 1.68
hlo_ffi / IPartOpt / cpu / PreRev 0.000024429 s 0.000015157460038608406 s 1.61
hlo_ffi / IPartOpt / cpu / PostRev 0.000024183 s 0.000014356760048030991 s 1.68
hlo_ffi / IPartOpt / cpu / BothRev 0.000024404 s 0.000014438420002989003 s 1.69
hlo_ffi / DefOpt / cpu / PreRev 0.000023952 s 0.000014575639988834157 s 1.64
hlo_ffi / DefOpt / cpu / PostRev 0.000024176 s 0.000014225739969333515 s 1.70
hlo_ffi / DefOpt / cpu / BothRev 0.000024614 s 0.000014324239991765353 s 1.72
hlo_ffi / IDefOpt / cpu / PreRev 0.00002392 s 0.00001521861997389351 s 1.57
hlo_ffi / IDefOpt / cpu / PostRev 0.000024017 s 0.000014588500034733444 s 1.65
hlo_ffi / IDefOpt / cpu / BothRev 0.000024283 s 0.000014252720002332352 s 1.70
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Primal 0.0009633138000936 s 0.0009321387999079 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Primal 0.0009354377996714 s 0.0009321677998741 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Primal 0.0009977226000046 s 0.0009918829998241 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Primal 0.0009997654000471 s 0.0009257460001208 s 1.08
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Primal 0.0009310389996244 s 0.0009386104000441 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Primal 0.0009784854000827 s 0.0010159395999835 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Primal 0.0009662222002589 s 0.0010079241998937 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Forward 0.0025441912001042 s 0.0022614704000261 s 1.13
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Forward 0.002362849399833 s 0.002405369400094 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Forward 0.0023368991998722 s 0.0022799684000347 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Forward 0.00228011900017 s 0.0022509202000946 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Forward 0.0022569457998542 s 0.0022421316000873 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Forward 0.00228969460004 s 0.0022940865999771 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Forward 0.0022838533999674 s 0.0022375612001269 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PreRev 0.0052257400002417 s 0.0056978842000717 s 0.92
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PostRev 0.0050017694000416 s 0.0063587201999325 s 0.79
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / BothRev 0.0054223017998083 s 0.0060367627999767 s 0.90
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / BothRev 0.0054415290000179 s 0.0063693950001834 s 0.85
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PreRev 0.0052041971997823 s 0.006341317999977 s 0.82
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PostRev 0.0052047602001039 s 0.0053775515999404 s 0.97
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / BothRev 0.0034103292000509 s 0.0055548254001223 s 0.61
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PreRev 0.0043516039997484 s 0.0055048000000169 s 0.79
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PostRev 0.0036082087999602 s 0.0061269304000234 s 0.59
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / BothRev 0.0034152631998949 s 0.0056070170000566 s 0.61
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PreRev 0.0050303773999985 s 0.006099837199963 s 0.82
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PostRev 0.005331957999806 s 0.0040929354000581 s 1.30
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / BothRev 0.0034391482000501 s 0.005288503400061 s 0.65
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PreRev 0.0050602249999428 s 0.0034711325999523 s 1.46
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PostRev 0.0033431271998779 s 0.0060114416000942 s 0.56
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / BothRev 0.0049138028000015 s 0.0035173248000319 s 1.40
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PreRev 0.0049412138001571 s 0.0034985649999725 s 1.41
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PostRev 0.0049455927999588 s 0.0056002957999226 s 0.88
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / BothRev 0.0050052344000505 s 0.0033613042000069 s 1.49
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / Primal 0.00028 s 0.000275552 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / Primal 0.000279681 s 0.00027424 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / Primal 0.0002866239999999 s 0.000290208 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / Primal 0.000279009 s 0.000275392 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / Primal 0.000279936 s 0.000276576 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / Primal 0.0002861759999999 s 0.000291777 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / Primal 0.00028752 s 0.000291072 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / Forward 0.000557759 s 0.0005613759999999 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / Forward 0.000539264 s 0.000542817 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / Forward 0.000558273 s 0.000560545 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / Forward 0.0005579199999999 s 0.000561408 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / Forward 0.000557664 s 0.000560993 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / Forward 0.000558272 s 0.000561793 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / Forward 0.000557889 s 0.000561408 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / PreRev 0.001025121 s 0.0010280979999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / PostRev 0.000985344 s 0.000987105 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / BothRev 0.0010252159999999 s 0.001029025 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / BothRev 0.000986945 s 0.000989153 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / PreRev 0.001012544 s 0.001013761 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / PostRev 0.001038241 s 0.001039041 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / BothRev 0.001012832 s 0.001013344 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / PreRev 0.001026753 s 0.001028897 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / PostRev 0.000974657 s 0.000978369 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / BothRev 0.001028289 s 0.001029601 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / PreRev 0.0010261119999999 s 0.0010275529999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / PostRev 0.000974816 s 0.0009758099999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / BothRev 0.001026368 s 0.001027169 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / PreRev 0.001022689 s 0.001024991 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / PostRev 0.0009595519999999 s 0.000964161 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / BothRev 0.001024064 s 0.001026305 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / PreRev 0.0010216 s 0.001024385 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / PostRev 0.001019296 s 0.001024928 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / BothRev 0.00102288 s 0.001023745 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / Primal 0.0001253605 s 0.0001232635 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / tpu / Primal 0.0001241637499999 s 0.0001233635 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / Primal 0.00015444525 s 0.00015185925 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / Primal 0.000131454 s 0.0001300629999999 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / Primal 0.000132552 s 0.000130404 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / Primal 0.00014509725 s 0.00014423625 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / Primal 0.00015328125 s 0.0001503662499999 s 1.02
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / Forward 0.000213223 s 0.0002128715 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / tpu / Forward 0.00026025475 s 0.000259804 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / Forward 0.0002203565 s 0.0002129352499999 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / Forward 0.00021139325 s 0.0002101505 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / Forward 0.00021562525 s 0.00021305025 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / Forward 0.000217957 s 0.0002101205 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / Forward 0.00021593125 s 0.0002127614999999 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / PreRev 0.0003550575 s 0.00035646575 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / PostRev 0.0002559755 s 0.0002569629999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / BothRev 0.00035494925 s 0.0003556017499999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / tpu / BothRev 0.0002569375 s 0.000257003 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / PreRev 0.00035501375 s 0.0003559415 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / PostRev 0.000291249 s 0.00029076975 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / BothRev 0.0003552435 s 0.00035559025 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / PreRev 0.0003555052499999 s 0.00035583025 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / PostRev 0.0002717335 s 0.00027396675 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / tpu / BothRev 0.000355911 s 0.00035568425 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / PreRev 0.000355241 s 0.0003553484999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / PostRev 0.0002721672499999 s 0.00027274175 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / BothRev 0.00035500625 s 0.0003554765 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / PreRev 0.00035802425 s 0.00035775425 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / PostRev 0.000283514 s 0.00028377575 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / tpu / BothRev 0.0003581587499999 s 0.0003579029999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / PreRev 0.00035740175 s 0.00035837225 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / PostRev 0.0003013454999999 s 0.0003013565 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / tpu / BothRev 0.00035751425 s 0.0003580575 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Primal 0.001675312 s 0.0009321387999079 s 1.80
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Primal 0.00163838 s 0.0009321677998741 s 1.76
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Primal 0.001685324 s 0.0009918829998241 s 1.70
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Primal 0.001720407 s 0.0009257460001208 s 1.86
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Primal 0.001526965 s 0.0009386104000441 s 1.63
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Primal 0.001577149 s 0.0010159395999835 s 1.55
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Primal 0.001653394 s 0.0010079241998937 s 1.64
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / Forward 0.0046035899999999 s 0.0022614704000261 s 2.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Forward 0.004974967 s 0.002405369400094 s 2.07
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Forward 0.0044497999999999 s 0.0022799684000347 s 1.95
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / Forward 0.004369557 s 0.0022509202000946 s 1.94
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / Forward 0.004605098 s 0.0022421316000873 s 2.05
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Forward 0.004576864 s 0.0022940865999771 s 2.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Forward 0.004776988 s 0.0022375612001269 s 2.13
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PreRev 0.008452354 s 0.0056978842000717 s 1.48
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / PostRev 0.009371574 s 0.0063587201999325 s 1.47
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / BothRev 0.008958051 s 0.0060367627999767 s 1.48
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / BothRev 0.008689821 s 0.0063693950001834 s 1.36
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PreRev 0.007950983 s 0.006341317999977 s 1.25
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PostRev 0.0082120099999999 s 0.0053775515999404 s 1.53
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / BothRev 0.007976406 s 0.0055548254001223 s 1.44
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PreRev 0.009316662 s 0.0055048000000169 s 1.69
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PostRev 0.009446977 s 0.0061269304000234 s 1.54
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / BothRev 0.008804573 s 0.0056070170000566 s 1.57
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PreRev 0.007813804 s 0.006099837199963 s 1.28
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PostRev 0.009584152 s 0.0040929354000581 s 2.34
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / BothRev 0.008186512 s 0.005288503400061 s 1.55
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PreRev 0.009127517 s 0.0034711325999523 s 2.63
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / PostRev 0.008085196 s 0.0060114416000942 s 1.34
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / BothRev 0.00920934 s 0.0035173248000319 s 2.62
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PreRev 0.008675559 s 0.0034985649999725 s 2.48
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PostRev 0.009043824 s 0.0056002957999226 s 1.61
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / BothRev 0.009451117 s 0.0033613042000069 s 2.81
scatter_sum / JaXPipe / cpu / Primal 0.000007649320068594533 s 0.00000753280001845269 s 1.02
scatter_sum / Jax / cpu / Primal 0.000007560620033473242 s 0.000008337480048794532 s 0.91
scatter_sum / HLOOpt / cpu / Primal 0.000008208980016206624 s 0.000007767980041535339 s 1.06
scatter_sum / PartOpt / cpu / Primal 0.000007294379920494975 s 0.00000764754005103896 s 0.95
scatter_sum / IPartOpt / cpu / Primal 0.000007376079993264284 s 0.000008083640050244867 s 0.91
scatter_sum / DefOpt / cpu / Primal 0.000007419220019073691 s 0.00000791293996371678 s 0.94
scatter_sum / IDefOpt / cpu / Primal 0.0000073918800444516815 s 0.000007435879997501615 s 0.99
scatter_sum / JaXPipe / cpu / Forward 0.000011385419929865747 s 0.000012307519973546732 s 0.93
scatter_sum / Jax / cpu / Forward 0.000011255339904892023 s 0.00001249038001333247 s 0.90
scatter_sum / HLOOpt / cpu / Forward 0.000011608619988692226 s 0.000012609119976332294 s 0.92
scatter_sum / PartOpt / cpu / Forward 0.000011278479960310503 s 0.00001196134005112981 s 0.94
scatter_sum / IPartOpt / cpu / Forward 0.000012079880016244714 s 0.000012671400008912317 s 0.95
scatter_sum / DefOpt / cpu / Forward 0.000011310000008961653 s 0.000012489519986047524 s 0.91
scatter_sum / IDefOpt / cpu / Forward 0.00001115807999667595 s 0.00001238691998878494 s 0.90
scatter_sum / JaXPipe / cpu / PreRev 0.000011653000092337606 s 0.000011618060007094757 s 1.00
scatter_sum / JaXPipe / cpu / PostRev 0.00001173426002424094 s 0.000012251960006324224 s 0.96
scatter_sum / JaXPipe / cpu / BothRev 0.000011775919992942364 s 0.00001236012001754716 s 0.95
scatter_sum / Jax / cpu / BothRev 0.000011863219988299534 s 0.00001187192001452786 s 1.00
scatter_sum / HLOOpt / cpu / PreRev 0.000011771680037782062 s 0.000012426240036802485 s 0.95
scatter_sum / HLOOpt / cpu / PostRev 0.000013234160069259816 s 0.000014015159995324212 s 0.94
scatter_sum / HLOOpt / cpu / BothRev 0.00001106186002289178 s 0.000011412799985919263 s 0.97
scatter_sum / PartOpt / cpu / PreRev 0.000011311420039419315 s 0.000011961099971813384 s 0.95
scatter_sum / PartOpt / cpu / PostRev 0.00001168115995824337 s 0.000011408780037527322 s 1.02
scatter_sum / PartOpt / cpu / BothRev 0.000012167879922344585 s 0.000012410079989422227 s 0.98
scatter_sum / IPartOpt / cpu / PreRev 0.0000112165999962599 s 0.000012210839995532295 s 0.92
scatter_sum / IPartOpt / cpu / PostRev 0.00001116396004363196 s 0.000011338740032442729 s 0.98
scatter_sum / IPartOpt / cpu / BothRev 0.000011276419991190778 s 0.000011471719999462948 s 0.98
scatter_sum / DefOpt / cpu / PreRev 0.000011668760016618762 s 0.000011868340016008003 s 0.98
scatter_sum / DefOpt / cpu / PostRev 0.000010984740019921449 s 0.000011632699988695096 s 0.94
scatter_sum / DefOpt / cpu / BothRev 0.000011654020036075962 s 0.000011750099965865956 s 0.99
scatter_sum / IDefOpt / cpu / PreRev 0.00001156021997303469 s 0.00001182487998448778 s 0.98
scatter_sum / IDefOpt / cpu / PostRev 0.000011028559983969898 s 0.000011796299995694426 s 0.93
scatter_sum / IDefOpt / cpu / BothRev 0.00001165523997769924 s 0.00001179338001747965 s 0.99
scatter_sum / JaXPipe / cuda / Primal 0.000009856 s 0.000010112 s 0.97
scatter_sum / Jax / cuda / Primal 0.000009696 s 0.000009824 s 0.99
scatter_sum / HLOOpt / cuda / Primal 0.000009952 s 0.000010112 s 0.98
scatter_sum / PartOpt / cuda / Primal 0.000009856 s 0.00001024 s 0.96
scatter_sum / IPartOpt / cuda / Primal 0.00001008 s 0.000010304 s 0.98
scatter_sum / DefOpt / cuda / Primal 0.000009344 s 0.000009791 s 0.95
scatter_sum / IDefOpt / cuda / Primal 0.000010112 s 0.000010239 s 0.99
scatter_sum / JaXPipe / cuda / Forward 0.000016768000000000003 s 0.000016896000000000002 s 0.99
scatter_sum / Jax / cuda / Forward 0.000016608 s 0.000017088 s 0.97
scatter_sum / HLOOpt / cuda / Forward 0.000016864 s 0.000017312 s 0.97
scatter_sum / PartOpt / cuda / Forward 0.00001632 s 0.000017056 s 0.96
scatter_sum / IPartOpt / cuda / Forward 0.000017024 s 0.000017152 s 0.99
scatter_sum / DefOpt / cuda / Forward 0.000016448999999999998 s 0.000016929 s 0.97
scatter_sum / IDefOpt / cuda / Forward 0.0000168 s 0.00002448 s 0.69
scatter_sum / JaXPipe / cuda / PreRev 0.000016352 s 0.000017503999999999997 s 0.93
scatter_sum / JaXPipe / cuda / PostRev 0.00001632 s 0.000016737 s 0.98
scatter_sum / JaXPipe / cuda / BothRev 0.000016896000000000002 s 0.000017152 s 0.99
scatter_sum / Jax / cuda / BothRev 0.000016864 s 0.000018753 s 0.90
scatter_sum / HLOOpt / cuda / PreRev 0.000016768999999999998 s 0.000017344 s 0.97
scatter_sum / HLOOpt / cuda / PostRev 0.000016864 s 0.000017312 s 0.97
scatter_sum / HLOOpt / cuda / BothRev 0.00001744 s 0.000017760000000000003 s 0.98
scatter_sum / PartOpt / cuda / PreRev 0.000016927999999999998 s 0.000018048 s 0.94
scatter_sum / PartOpt / cuda / PostRev 0.000016352 s 0.000017760000000000003 s 0.92
scatter_sum / PartOpt / cuda / BothRev 0.000018752000000000003 s 0.000018016 s 1.04
scatter_sum / IPartOpt / cuda / PreRev 0.000016704 s 0.000017888999999999998 s 0.93
scatter_sum / IPartOpt / cuda / PostRev 0.000015808 s 0.000016670999999999997 s 0.95
scatter_sum / IPartOpt / cuda / BothRev 0.000016864 s 0.000017888999999999998 s 0.94
scatter_sum / DefOpt / cuda / PreRev 0.000016927999999999998 s 0.000017152 s 0.99
scatter_sum / DefOpt / cuda / PostRev 0.000016864 s 0.000017152 s 0.98
scatter_sum / DefOpt / cuda / BothRev 0.000018815 s 0.000017216 s 1.09
scatter_sum / IDefOpt / cuda / PreRev 0.000017088 s 0.00001696 s 1.01
scatter_sum / IDefOpt / cuda / PostRev 0.000016992 s 0.000017247999999999998 s 0.99
scatter_sum / IDefOpt / cuda / BothRev 0.000016927999999999998 s 0.000017345 s 0.98
scatter_sum / JaXPipe / tpu / Primal 0.0000013505000000000002 s 0.000001343275 s 1.01
scatter_sum / Jax / tpu / Primal 0.0000014045500000000002 s 0.000001404425 s 1.00
scatter_sum / HLOOpt / tpu / Primal 0.000001350925 s 0.00000134315 s 1.01
scatter_sum / PartOpt / tpu / Primal 0.00000140505 s 0.000001404425 s 1.00
scatter_sum / IPartOpt / tpu / Primal 0.000001350675 s 0.0000013434250000000002 s 1.01
scatter_sum / DefOpt / tpu / Primal 0.000001404625 s 0.00000140475 s 1.00
scatter_sum / IDefOpt / tpu / Primal 0.000001350925 s 0.0000013432 s 1.01
scatter_sum / JaXPipe / tpu / Forward 0.0000027028 s 0.000002704675 s 1.00
scatter_sum / Jax / tpu / Forward 0.000002736375 s 0.00000271835 s 1.01
scatter_sum / HLOOpt / tpu / Forward 0.00000270805 s 0.0000027015 s 1.00
scatter_sum / PartOpt / tpu / Forward 0.0000026969 s 0.00000268595 s 1.00
scatter_sum / IPartOpt / tpu / Forward 0.000002702475 s 0.0000027071 s 1.00
scatter_sum / DefOpt / tpu / Forward 0.000002696375 s 0.0000026979 s 1.00
scatter_sum / IDefOpt / tpu / Forward 0.0000027018999999999995 s 0.000002715275 s 1.00
scatter_sum / JaXPipe / tpu / PreRev 0.000002691525 s 0.0000026796750000000003 s 1.00
scatter_sum / JaXPipe / tpu / PostRev 0.0000026933 s 0.000002681375 s 1.00
scatter_sum / JaXPipe / tpu / BothRev 0.00000271145 s 0.0000026992500000000005 s 1.00
scatter_sum / Jax / tpu / BothRev 0.000002742225 s 0.000002735225 s 1.00
scatter_sum / HLOOpt / tpu / PreRev 0.000002707 s 0.000002694525 s 1.00
scatter_sum / HLOOpt / tpu / PostRev 0.000002742225 s 0.000002740225 s 1.00
scatter_sum / HLOOpt / tpu / BothRev 0.0000027053 s 0.0000026929750000000003 s 1.00
scatter_sum / PartOpt / tpu / PreRev 0.0000027531 s 0.000002751425 s 1.00
scatter_sum / PartOpt / tpu / PostRev 0.000002704725 s 0.0000026935 s 1.00
scatter_sum / PartOpt / tpu / BothRev 0.000002745025 s 0.0000027455 s 1.00
scatter_sum / IPartOpt / tpu / PreRev 0.0000027098 s 0.0000026997500000000003 s 1.00
scatter_sum / IPartOpt / tpu / PostRev 0.0000027415 s 0.00000273565 s 1.00
scatter_sum / IPartOpt / tpu / BothRev 0.0000027101750000000003 s 0.000002711575 s 1.00
scatter_sum / DefOpt / tpu / PreRev 0.0000027402 s 0.00000274295 s 1.00
scatter_sum / DefOpt / tpu / PostRev 0.000002702525 s 0.000002690275 s 1.00
scatter_sum / DefOpt / tpu / BothRev 0.000002740225 s 0.0000027378750000000004 s 1.00
scatter_sum / IDefOpt / tpu / PreRev 0.0000027035 s 0.0000026918750000000004 s 1.00
scatter_sum / IDefOpt / tpu / PostRev 0.000002743525 s 0.00000273465 s 1.00
scatter_sum / IDefOpt / tpu / BothRev 0.000002706575 s 0.000002694975 s 1.00
scatter_sum / JaXPipe / cpu / Primal 0.000015869 s 0.00000753280001845269 s 2.11
scatter_sum / Jax / cpu / Primal 0.000015787 s 0.000008337480048794532 s 1.89
scatter_sum / HLOOpt / cpu / Primal 0.000015601000000000002 s 0.000007767980041535339 s 2.01
scatter_sum / PartOpt / cpu / Primal 0.000015431000000000002 s 0.00000764754005103896 s 2.02
scatter_sum / IPartOpt / cpu / Primal 0.000015694 s 0.000008083640050244867 s 1.94
scatter_sum / DefOpt / cpu / Primal 0.000015583 s 0.00000791293996371678 s 1.97
scatter_sum / IDefOpt / cpu / Primal 0.000015615 s 0.000007435879997501615 s 2.10
scatter_sum / JaXPipe / cpu / Forward 0.000037718 s 0.000012307519973546732 s 3.06
scatter_sum / Jax / cpu / Forward 0.00002242 s 0.00001249038001333247 s 1.79
scatter_sum / HLOOpt / cpu / Forward 0.000021754 s 0.000012609119976332294 s 1.73
scatter_sum / PartOpt / cpu / Forward 0.000022192 s 0.00001196134005112981 s 1.86
scatter_sum / IPartOpt / cpu / Forward 0.000022357 s 0.000012671400008912317 s 1.76
scatter_sum / DefOpt / cpu / Forward 0.000022654 s 0.000012489519986047524 s 1.81
scatter_sum / IDefOpt / cpu / Forward 0.000022639 s 0.00001238691998878494 s 1.83
scatter_sum / JaXPipe / cpu / PreRev 0.000022778 s 0.000011618060007094757 s 1.96
scatter_sum / JaXPipe / cpu / PostRev 0.000022548 s 0.000012251960006324224 s 1.84
scatter_sum / JaXPipe / cpu / BothRev 0.000022331 s 0.00001236012001754716 s 1.81
scatter_sum / Jax / cpu / BothRev 0.000022811 s 0.00001187192001452786 s 1.92
scatter_sum / HLOOpt / cpu / PreRev 0.000022693 s 0.000012426240036802485 s 1.83
scatter_sum / HLOOpt / cpu / PostRev 0.00002209 s 0.000014015159995324212 s 1.58
scatter_sum / HLOOpt / cpu / BothRev 0.000022976 s 0.000011412799985919263 s 2.01
scatter_sum / PartOpt / cpu / PreRev 0.000022576 s 0.000011961099971813384 s 1.89
scatter_sum / PartOpt / cpu / PostRev 0.000022462 s 0.000011408780037527322 s 1.97
scatter_sum / PartOpt / cpu / BothRev 0.000022269 s 0.000012410079989422227 s 1.79
scatter_sum / IPartOpt / cpu / PreRev 0.000022762 s 0.000012210839995532295 s 1.86
scatter_sum / IPartOpt / cpu / PostRev 0.000022069 s 0.000011338740032442729 s 1.95
scatter_sum / IPartOpt / cpu / BothRev 0.000022335 s 0.000011471719999462948 s 1.95
scatter_sum / DefOpt / cpu / PreRev 0.000022762 s 0.000011868340016008003 s 1.92
scatter_sum / DefOpt / cpu / PostRev 0.000022746 s 0.000011632699988695096 s 1.96
scatter_sum / DefOpt / cpu / BothRev 0.000022105 s 0.000011750099965865956 s 1.88
scatter_sum / IDefOpt / cpu / PreRev 0.000022481 s 0.00001182487998448778 s 1.90
scatter_sum / IDefOpt / cpu / PostRev 0.000022496 s 0.000011796299995694426 s 1.91
scatter_sum / IDefOpt / cpu / BothRev 0.000022321 s 0.00001179338001747965 s 1.89
slicing / JaXPipe / cpu / Primal 0.00000619050004388555 s 0.000006597080018764246 s 0.94
slicing / Jax / cpu / Primal 0.000006004019960528239 s 0.000006177960012792028 s 0.97
slicing / HLOOpt / cpu / Primal 0.000006901919932715827 s 0.000006341799980873475 s 1.09
slicing / PartOpt / cpu / Primal 0.000006827599954704056 s 0.000006137679984021815 s 1.11
slicing / IPartOpt / cpu / Primal 0.000006270279991440475 s 0.000006653579976045876 s 0.94
slicing / DefOpt / cpu / Primal 0.000006089239941502456 s 0.000006134799996289075 s 0.99
slicing / IDefOpt / cpu / Primal 0.000006721619884046959 s 0.000006215719995452673 s 1.08
slicing / JaXPipe / cpu / Forward 0.00000952321994191152 s 0.00000958000001446635 s 0.99
slicing / Jax / cpu / Forward 0.000009676440022303725 s 0.00000932574000216846 s 1.04
slicing / HLOOpt / cpu / Forward 0.00000920937998671434 s 0.000009797440025067772 s 0.94
slicing / PartOpt / cpu / Forward 0.000008975460023066262 s 0.00000952871996560134 s 0.94
slicing / IPartOpt / cpu / Forward 0.00000926763994357316 s 0.00001003165997644828 s 0.92
slicing / DefOpt / cpu / Forward 0.000009126660097535933 s 0.000009296859998357831 s 0.98
slicing / IDefOpt / cpu / Forward 0.000009718319943203825 s 0.000009523819971946069 s 1.02
slicing / JaXPipe / cpu / PreRev 0.00000991310005701962 s 0.000010216460013907636 s 0.97
slicing / JaXPipe / cpu / PostRev 0.000009965959907276557 s 0.000009857619970716768 s 1.01
slicing / JaXPipe / cpu / BothRev 0.000009847859964793314 s 0.000010476039997229235 s 0.94
slicing / Jax / cpu / BothRev 0.000009636460035835623 s 0.000009995319987865516 s 0.96
slicing / HLOOpt / cpu / PreRev 0.000010174620019824942 s 0.000010662620015864376 s 0.95
slicing / HLOOpt / cpu / PostRev 0.000011779460019170074 s 0.0000123771199832845 s 0.95
slicing / HLOOpt / cpu / BothRev 0.000009516400059510488 s 0.0000095578600121371 s 1.00
slicing / PartOpt / cpu / PreRev 0.000009597340031177738 s 0.000009616839970476575 s 1.00
slicing / PartOpt / cpu / PostRev 0.000009548079979140312 s 0.0000105017599707935 s 0.91
slicing / PartOpt / cpu / BothRev 0.000010063640002044848 s 0.000010213999994448386 s 0.99
slicing / IPartOpt / cpu / PreRev 0.000009297519882238703 s 0.000010010960013460136 s 0.93
slicing / IPartOpt / cpu / PostRev 0.000010622199952194932 s 0.000010388719992988626 s 1.02
slicing / IPartOpt / cpu / BothRev 0.000010289940128132004 s 0.000009856840006250422 s 1.04
slicing / DefOpt / cpu / PreRev 0.000009857339955487987 s 0.000009915399969031567 s 0.99
slicing / DefOpt / cpu / PostRev 0.000009510540057817708 s 0.000010333379996154693 s 0.92
slicing / DefOpt / cpu / BothRev 0.00001048715994329541 s 0.000010391340001660863 s 1.01
slicing / IDefOpt / cpu / PreRev 0.00000962163998337928 s 0.00000994755997453467 s 0.97
slicing / IDefOpt / cpu / PostRev 0.000009851760005403775 s 0.00000992942003904318 s 0.99
slicing / IDefOpt / cpu / BothRev 0.000009526520098006583 s 0.000009599200056982228 s 0.99
slicing / JaXPipe / cuda / Primal 0.000001887 s 0.000001887 s 1
slicing / Jax / cuda / Primal 0.000001887 s 0.000001887 s 1
slicing / HLOOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
slicing / PartOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
slicing / IPartOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
slicing / DefOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
slicing / IDefOpt / cuda / Primal 0.000001887 s 0.000001887 s 1
slicing / JaXPipe / cuda / Forward 0.000009632 s 0.000009856 s 0.98
slicing / Jax / cuda / Forward 0.000009856 s 0.000010175 s 0.97
slicing / HLOOpt / cuda / Forward 0.000010016 s 0.000010048 s 1.00
slicing / PartOpt / cuda / Forward 0.000009376 s 0.000010176 s 0.92
slicing / IPartOpt / cuda / Forward 0.000010176 s 0.000009856 s 1.03
slicing / DefOpt / cuda / Forward 0.000009856 s 0.00000992 s 0.99
slicing / IDefOpt / cuda / Forward 0.00001008 s 0.000010145 s 0.99
slicing / JaXPipe / cuda / PreRev 0.000009569 s 0.000009856 s 0.97
slicing / JaXPipe / cuda / PostRev 0.000009824 s 0.000010431 s 0.94
slicing / JaXPipe / cuda / BothRev 0.0000096 s 0.000010112 s 0.95
slicing / Jax / cuda / BothRev 0.000010015 s 0.00001024 s 0.98
slicing / HLOOpt / cuda / PreRev 0.000009024 s 0.000009888 s 0.91
slicing / HLOOpt / cuda / PostRev 0.00000944 s 0.000009664 s 0.98
slicing / HLOOpt / cuda / BothRev 0.000009983 s 0.000009952 s 1.00
slicing / PartOpt / cuda / PreRev 0.00000976 s 0.00001008 s 0.97
slicing / PartOpt / cuda / PostRev 0.000009728 s 0.000009792 s 0.99
slicing / PartOpt / cuda / BothRev 0.000009504 s 0.000009824 s 0.97
slicing / IPartOpt / cuda / PreRev 0.000009408 s 0.00000976 s 0.96
slicing / IPartOpt / cuda / PostRev 0.000009504 s 0.000009824 s 0.97
slicing / IPartOpt / cuda / BothRev 0.000009696 s 0.000009825 s 0.99
slicing / DefOpt / cuda / PreRev 0.000009792 s 0.000009856 s 0.99
slicing / DefOpt / cuda / PostRev 0.000009984 s 0.000010048 s 0.99
slicing / DefOpt / cuda / BothRev 0.000009504 s 0.000009791 s 0.97
slicing / IDefOpt / cuda / PreRev 0.000009824 s 0.00001008 s 0.97
slicing / IDefOpt / cuda / PostRev 0.000009344 s 0.000010111 s 0.92
slicing / IDefOpt / cuda / BothRev 0.000009472 s 0.000010048 s 0.94
slicing / JaXPipe / tpu / Primal 0.0000010268 s 9.744e-7 s 1.05
slicing / Jax / tpu / Primal 9.734e-7 s 9.7205e-7 s 1.00
slicing / HLOOpt / tpu / Primal 0.0000010256 s 9.65325e-7 s 1.06
slicing / PartOpt / tpu / Primal 9.82225e-7 s 9.68125e-7 s 1.01
slicing / IPartOpt / tpu / Primal 0.0000010280000000000002 s 9.672999999999998e-7 s 1.06
slicing / DefOpt / tpu / Primal 9.693e-7 s 9.6845e-7 s 1.00
slicing / IDefOpt / tpu / Primal 0.000001024425 s 9.64425e-7 s 1.06
slicing / JaXPipe / tpu / Forward 0.000001415325 s 0.000001409775 s 1.00
slicing / Jax / tpu / Forward 0.00000147905 s 0.000001425125 s 1.04
slicing / HLOOpt / tpu / Forward 0.00000152125 s 0.0000015266 s 1.00
slicing / PartOpt / tpu / Forward 0.00000149995 s 0.000001438725 s 1.04
slicing / IPartOpt / tpu / Forward 0.000001521825 s 0.0000015182 s 1.00
slicing / DefOpt / tpu / Forward 0.0000014987 s 0.000001438125 s 1.04
slicing / IDefOpt / tpu / Forward 0.000001517475 s 0.00000151785 s 1.00
slicing / JaXPipe / tpu / PreRev 0.00000256315 s 0.000002385175 s 1.07
slicing / JaXPipe / tpu / PostRev 0.000002509575 s 0.0000025197 s 1.00
slicing / JaXPipe / tpu / BothRev 0.000002587475 s 0.0000024034500000000003 s 1.08
slicing / Jax / tpu / BothRev 0.00000253735 s 0.000002536525 s 1.00
slicing / HLOOpt / tpu / PreRev 0.0000025802000000000003 s 0.000002401 s 1.07
slicing / HLOOpt / tpu / PostRev 0.000002530475 s 0.0000025440750000000003 s 0.99
slicing / HLOOpt / tpu / BothRev 0.000002595825 s 0.00000239565 s 1.08
slicing / PartOpt / tpu / PreRev 0.000002543025 s 0.000002539375 s 1.00
slicing / PartOpt / tpu / PostRev 0.000002589675 s 0.0000023967250000000003 s 1.08
slicing / PartOpt / tpu / BothRev 0.000002536775 s 0.0000025396 s 1.00
slicing / IPartOpt / tpu / PreRev 0.00000257905 s 0.000002407025 s 1.07
slicing / IPartOpt / tpu / PostRev 0.0000025337 s 0.000002555175 s 0.99
slicing / IPartOpt / tpu / BothRev 0.00000259605 s 0.000002390125 s 1.09
slicing / DefOpt / tpu / PreRev 0.000002542325 s 0.000002549 s 1.00
slicing / DefOpt / tpu / PostRev 0.000002589425 s 0.00000240395 s 1.08
slicing / DefOpt / tpu / BothRev 0.000002536275 s 0.0000025439 s 1.00
slicing / IDefOpt / tpu / PreRev 0.00000258065 s 0.00000240165 s 1.07
slicing / IDefOpt / tpu / PostRev 0.000002537425 s 0.0000025458750000000004 s 1.00
slicing / IDefOpt / tpu / BothRev 0.000002590225 s 0.0000024013250000000003 s 1.08
slicing / JaXPipe / cpu / Primal 0.000012553 s 0.000006597080018764246 s 1.90
slicing / Jax / cpu / Primal 0.000012562 s 0.000006177960012792028 s 2.03
slicing / HLOOpt / cpu / Primal 0.000012297 s 0.000006341799980873475 s 1.94
slicing / PartOpt / cpu / Primal 0.000012568 s 0.000006137679984021815 s 2.05
slicing / IPartOpt / cpu / Primal 0.000012447 s 0.000006653579976045876 s 1.87
slicing / DefOpt / cpu / Primal 0.000012432 s 0.000006134799996289075 s 2.03
slicing / IDefOpt / cpu / Primal 0.000012551 s 0.000006215719995452673 s 2.02
slicing / JaXPipe / cpu / Forward 0.000016975 s 0.00000958000001446635 s 1.77
slicing / Jax / cpu / Forward 0.000016902000000000002 s 0.00000932574000216846 s 1.81
slicing / HLOOpt / cpu / Forward 0.00001684 s 0.000009797440025067772 s 1.72
slicing / PartOpt / cpu / Forward 0.000016632999999999998 s 0.00000952871996560134 s 1.75
slicing / IPartOpt / cpu / Forward 0.000016817 s 0.00001003165997644828 s 1.68
slicing / DefOpt / cpu / Forward 0.00001651 s 0.000009296859998357831 s 1.78
slicing / IDefOpt / cpu / Forward 0.000016565 s 0.000009523819971946069 s 1.74
slicing / JaXPipe / cpu / PreRev 0.000017666 s 0.000010216460013907636 s 1.73
slicing / JaXPipe / cpu / PostRev 0.000023502 s 0.000009857619970716768 s 2.38
slicing / JaXPipe / cpu / BothRev 0.000017477 s 0.000010476039997229235 s 1.67
slicing / Jax / cpu / BothRev 0.000017422 s 0.000009995319987865516 s 1.74
slicing / HLOOpt / cpu / PreRev 0.000017301 s 0.000010662620015864376 s 1.62
slicing / HLOOpt / cpu / PostRev 0.000017257 s 0.0000123771199832845 s 1.39
slicing / HLOOpt / cpu / BothRev 0.000017182 s 0.0000095578600121371 s 1.80
slicing / PartOpt / cpu / PreRev 0.000017312 s 0.000009616839970476575 s 1.80
slicing / PartOpt / cpu / PostRev 0.000017419 s 0.0000105017599707935 s 1.66
slicing / PartOpt / cpu / BothRev 0.000017638 s 0.000010213999994448386 s 1.73
slicing / IPartOpt / cpu / PreRev 0.000017503999999999997 s 0.000010010960013460136 s 1.75
slicing / IPartOpt / cpu / PostRev 0.000017239000000000002 s 0.000010388719992988626 s 1.66
slicing / IPartOpt / cpu / BothRev 0.000016989 s 0.000009856840006250422 s 1.72
slicing / DefOpt / cpu / PreRev 0.000017378999999999997 s 0.000009915399969031567 s 1.75
slicing / DefOpt / cpu / PostRev 0.000017473 s 0.000010333379996154693 s 1.69
slicing / DefOpt / cpu / BothRev 0.000017292 s 0.000010391340001660863 s 1.66
slicing / IDefOpt / cpu / PreRev 0.000017478 s 0.00000994755997453467 s 1.76
slicing / IDefOpt / cpu / PostRev 0.000017049999999999998 s 0.00000992942003904318 s 1.72
slicing / IDefOpt / cpu / BothRev 0.000017497 s 0.000009599200056982228 s 1.82
sum / JaXPipe / cpu / Primal 0.000008003599923540605 s 0.000007966920029502945 s 1.00
sum / Jax / cpu / Primal 0.000007328579940804048 s 0.000007496000016544712 s 0.98
sum / HLOOpt / cpu / Primal 0.000008093160031421576 s 0.0000075576799827103965 s 1.07
sum / PartOpt / cpu / Primal 0.000007591720077471109 s 0.000007438460024786764 s 1.02
sum / IPartOpt / cpu / Primal 0.000007544600030087167 s 0.000008263000008810195 s 0.91
sum / DefOpt / cpu / Primal 0.000007415800100716297 s 0.000007974939999257913 s 0.93
sum / IDefOpt / cpu / Primal 0.000007740699984424282 s 0.000008394519963985659 s 0.92
sum / JaXPipe / cpu / Forward 0.00001128392001191969 s 0.00001140052003393066 s 0.99
sum / Jax / cpu / Forward 0.00001135973994678352 s 0.000011591580005188009 s 0.98
sum / HLOOpt / cpu / Forward 0.000011001460043189582 s 0.000011855700004161916 s 0.93
sum / PartOpt / cpu / Forward 0.00001073477997124428 s 0.000011221519962418824 s 0.96
sum / IPartOpt / cpu / Forward 0.00001172653996036388 s 0.000011349360020176392 s 1.03
sum / DefOpt / cpu / Forward 0.000011067680025007576 s 0.00001117800003157754 s 0.99
sum / IDefOpt / cpu / Forward 0.00001126471997849876 s 0.000011447939978097565 s 0.98
sum / JaXPipe / cpu / PreRev 0.000010479899883648611 s 0.000010888779997912936 s 0.96
sum / JaXPipe / cpu / PostRev 0.000010555799963185564 s 0.000010675600033209777 s 0.99
sum / JaXPipe / cpu / BothRev 0.000010519340103201102 s 0.00001096442001653486 s 0.96
sum / Jax / cpu / BothRev 0.000010460980029165512 s 0.000011158399984196876 s 0.94
sum / HLOOpt / cpu / PreRev 0.000011273039963271005 s 0.00001121798000895069 s 1.00
sum / HLOOpt / cpu / PostRev 0.000012629900047613774 s 0.000013005540022277274 s 0.97
sum / HLOOpt / cpu / BothRev 0.00001097227999707684 s 0.000011337859968989505 s 0.97
sum / PartOpt / cpu / PreRev 0.000011059199950977928 s 0.000011103740016551455 s 1.00
sum / PartOpt / cpu / PostRev 0.000010810979947564192 s 0.000010653500030457508 s 1.01
sum / PartOpt / cpu / BothRev 0.000011064320060540922 s 0.00001118933999350702 s 0.99
sum / IPartOpt / cpu / PreRev 0.000010672160005924524 s 0.000011035779980375082 s 0.97
sum / IPartOpt / cpu / PostRev 0.000010322299967810975 s 0.000011056240000471008 s 0.93
sum / IPartOpt / cpu / BothRev 0.000010423979929328198 s 0.000010846880031749608 s 0.96
sum / DefOpt / cpu / PreRev 0.000010416639925097116 s 0.000011112760003015865 s 0.94
sum / DefOpt / cpu / PostRev 0.000010426740045659244 s 0.000010988719977831352 s 0.95
sum / DefOpt / cpu / BothRev 0.000010835659977601609 s 0.000010721840008045548 s 1.01
sum / IDefOpt / cpu / PreRev 0.000010771280012704664 s 0.00001134413999352546 s 0.95
sum / IDefOpt / cpu / PostRev 0.000010518999915802852 s 0.000010876399992412189 s 0.97
sum / IDefOpt / cpu / BothRev 0.00001033693994031637 s 0.000011014359934051754 s 0.94
sum / JaXPipe / cuda / Primal 0.000002047 s 0.000002047 s 1
sum / Jax / cuda / Primal 0.000002048 s 0.000002048 s 1
sum / HLOOpt / cuda / Primal 0.000002048 s 0.000002047 s 1.00
sum / PartOpt / cuda / Primal 0.000002048 s 0.000002047 s 1.00
sum / IPartOpt / cuda / Primal 0.000002048 s 0.000002047 s 1.00
sum / DefOpt / cuda / Primal 0.000002047 s 0.000002048 s 1.00
sum / IDefOpt / cuda / Primal 0.000002047 s 0.000002047 s 1
sum / JaXPipe / cuda / Forward 0.000009792 s 0.000015263999999999998 s 0.64
sum / Jax / cuda / Forward 0.000010175 s 0.000010464 s 0.97
sum / HLOOpt / cuda / Forward 0.000009952 s 0.0000104 s 0.96
sum / PartOpt / cuda / Forward 0.000010304 s 0.000010304 s 1
sum / IPartOpt / cuda / Forward 0.000009792 s 0.000010272 s 0.95
sum / DefOpt / cuda / Forward 0.000009793 s 0.000010208 s 0.96
sum / IDefOpt / cuda / Forward 0.000010112 s 0.000010336 s 0.98
sum / JaXPipe / cuda / PreRev 0.000009504 s 0.000010048 s 0.95
sum / JaXPipe / cuda / PostRev 0.00000944 s 0.00001008 s 0.94
sum / JaXPipe / cuda / BothRev 0.000010144 s 0.000014112 s 0.72
sum / Jax / cuda / BothRev 0.000010464 s 0.0000096 s 1.09
sum / HLOOpt / cuda / PreRev 0.000009024 s 0.000009952 s 0.91
sum / HLOOpt / cuda / PostRev 0.000009376 s 0.000009633 s 0.97
sum / HLOOpt / cuda / BothRev 0.00000944 s 0.00001008 s 0.94
sum / PartOpt / cuda / PreRev 0.000009312000000000002 s 0.000010048 s 0.93
sum / PartOpt / cuda / PostRev 0.000009408 s 0.000009313 s 1.01
sum / PartOpt / cuda / BothRev 0.000009408 s 0.00000944 s 1.00
sum / IPartOpt / cuda / PreRev 0.000009631 s 0.000010047 s 0.96
sum / IPartOpt / cuda / PostRev 0.00000944 s 0.000009376 s 1.01
sum / IPartOpt / cuda / BothRev 0.000009632 s 0.000009984 s 0.96
sum / DefOpt / cuda / PreRev 0.00000976 s 0.00001008 s 0.97
sum / DefOpt / cuda / PostRev 0.000009727 s 0.00001008 s 0.96
sum / DefOpt / cuda / BothRev 0.00000928 s 0.000014433 s 0.64
sum / IDefOpt / cuda / PreRev 0.000009632 s 0.000010048 s 0.96
sum / IDefOpt / cuda / PostRev 0.000010368 s 0.000010112 s 1.03
sum / IDefOpt / cuda / BothRev 0.000010688 s 0.0000096 s 1.11
sum / JaXPipe / tpu / Primal 5.099999999999999e-7 s 5.103250000000001e-7 s 1.00
sum / Jax / tpu / Primal 5.47275e-7 s 5.47525e-7 s 1.00
sum / HLOOpt / tpu / Primal 5.101e-7 s 5.1055e-7 s 1.00
sum / PartOpt / tpu / Primal 5.47125e-7 s 5.47175e-7 s 1.00
sum / IPartOpt / tpu / Primal 5.1015e-7 s 5.10875e-7 s 1.00
sum / DefOpt / tpu / Primal 5.47175e-7 s 5.472250000000001e-7 s 1.00
sum / IDefOpt / tpu / Primal 5.104750000000001e-7 s 5.10825e-7 s 1.00
sum / JaXPipe / tpu / Forward 0.0000015518499999999998 s 0.000001557325 s 1.00
sum / Jax / tpu / Forward 0.00000149815 s 0.0000014961749999999997 s 1.00
sum / HLOOpt / tpu / Forward 0.0000015334 s 0.000001532275 s 1.00
sum / PartOpt / tpu / Forward 0.0000014979749999999998 s 0.000001497125 s 1.00
sum / IPartOpt / tpu / Forward 0.0000015304750000000002 s 0.000001538 s 1.00
sum / DefOpt / tpu / Forward 0.0000015007 s 0.0000015041 s 1.00
sum / IDefOpt / tpu / Forward 0.0000015336 s 0.0000015345 s 1.00
sum / JaXPipe / tpu / PreRev 0.000001051825 s 0.000001003375 s 1.05
sum / JaXPipe / tpu / PostRev 0.0000011009 s 0.0000010455 s 1.05
sum / JaXPipe / tpu / BothRev 0.0000010475 s 9.98625e-7 s 1.05
sum / Jax / tpu / BothRev 0.000001083475 s 0.00000104185 s 1.04
sum / HLOOpt / tpu / PreRev 0.000001047425 s 0.00000101055 s 1.04
sum / HLOOpt / tpu / PostRev 0.0000010908249999999998 s 0.0000010539 s 1.04
sum / HLOOpt / tpu / BothRev 0.00000105115 s 0.000001004325 s 1.05
sum / PartOpt / tpu / PreRev 0.00000108685 s 0.0000010631 s 1.02
sum / PartOpt / tpu / PostRev 0.00000104785 s 0.00000101205 s 1.04
sum / PartOpt / tpu / BothRev 0.00000109045 s 0.000001043275 s 1.05
sum / IPartOpt / tpu / PreRev 0.00000104835 s 0.000001023525 s 1.02
sum / IPartOpt / tpu / PostRev 0.00000108505 s 0.000001059325 s 1.02
sum / IPartOpt / tpu / BothRev 0.0000010545 s 0.0000010066999999999998 s 1.05
sum / DefOpt / tpu / PreRev 0.0000010852000000000002 s 0.0000010547 s 1.03
sum / DefOpt / tpu / PostRev 0.0000010485750000000002 s 0.000001002725 s 1.05
sum / DefOpt / tpu / BothRev 0.00000109165 s 0.0000010368 s 1.05
sum / IDefOpt / tpu / PreRev 0.000001056 s 0.000001004125 s 1.05
sum / IDefOpt / tpu / PostRev 0.0000010897 s 0.000001046575 s 1.04
sum / IDefOpt / tpu / BothRev 0.000001048825 s 0.000001000425 s 1.05
sum / JaXPipe / cpu / Primal 0.000014615 s 0.000007966920029502945 s 1.83
sum / Jax / cpu / Primal 0.000014707 s 0.000007496000016544712 s 1.96
sum / HLOOpt / cpu / Primal 0.000014597 s 0.0000075576799827103965 s 1.93
sum / PartOpt / cpu / Primal 0.000014795 s 0.000007438460024786764 s 1.99
sum / IPartOpt / cpu / Primal 0.000014795 s 0.000008263000008810195 s 1.79
sum / DefOpt / cpu / Primal 0.000014582 s 0.000007974939999257913 s 1.83
sum / IDefOpt / cpu / Primal 0.000014197 s 0.000008394519963985659 s 1.69
sum / JaXPipe / cpu / Forward 0.000020176 s 0.00001140052003393066 s 1.77
sum / Jax / cpu / Forward 0.000020003 s 0.000011591580005188009 s 1.73
sum / HLOOpt / cpu / Forward 0.000019882 s 0.000011855700004161916 s 1.68
sum / PartOpt / cpu / Forward 0.000020205 s 0.000011221519962418824 s 1.80
sum / IPartOpt / cpu / Forward 0.00001961 s 0.000011349360020176392 s 1.73
sum / DefOpt / cpu / Forward 0.000020176 s 0.00001117800003157754 s 1.80
sum / IDefOpt / cpu / Forward 0.000019658 s 0.000011447939978097565 s 1.72
sum / JaXPipe / cpu / PreRev 0.000019043 s 0.000010888779997912936 s 1.75
sum / JaXPipe / cpu / PostRev 0.000018795 s 0.000010675600033209777 s 1.76
sum / JaXPipe / cpu / BothRev 0.000018705 s 0.00001096442001653486 s 1.71
sum / Jax / cpu / BothRev 0.000018941 s 0.000011158399984196876 s 1.70
sum / HLOOpt / cpu / PreRev 0.00001847 s 0.00001121798000895069 s 1.65
sum / HLOOpt / cpu / PostRev 0.000019137 s 0.000013005540022277274 s 1.47
sum / HLOOpt / cpu / BothRev 0.000019034 s 0.000011337859968989505 s 1.68
sum / PartOpt / cpu / PreRev 0.000018842 s 0.000011103740016551455 s 1.70
sum / PartOpt / cpu / PostRev 0.000018849 s 0.000010653500030457508 s 1.77
sum / PartOpt / cpu / BothRev 0.000018993 s 0.00001118933999350702 s 1.70
sum / IPartOpt / cpu / PreRev 0.000018916 s 0.000011035779980375082 s 1.71
sum / IPartOpt / cpu / PostRev 0.000019041 s 0.000011056240000471008 s 1.72
sum / IPartOpt / cpu / BothRev 0.000018610000000000003 s 0.000010846880031749608 s 1.72
sum / DefOpt / cpu / PreRev 0.000019031 s 0.000011112760003015865 s 1.71
sum / DefOpt / cpu / PostRev 0.000018718 s 0.000010988719977831352 s 1.70
sum / DefOpt / cpu / BothRev 0.000018714 s 0.000010721840008045548 s 1.75
sum / IDefOpt / cpu / PreRev 0.000018905000000000003 s 0.00001134413999352546 s 1.67
sum / IDefOpt / cpu / PostRev 0.000018815 s 0.000010876399992412189 s 1.73
sum / IDefOpt / cpu / BothRev 0.000019063 s 0.000011014359934051754 s 1.73
value_and_grad / JaXPipe / cpu / Primal 0.00001395937997585861 s 0.00001462693999201292 s 0.95
value_and_grad / Jax / cpu / Primal 0.00001315113993769046 s 0.000013887280019844183 s 0.95
value_and_grad / HLOOpt / cpu / Primal 0.000013560220068029594 s 0.000013906800022596144 s 0.98
value_and_grad / PartOpt / cpu / Primal 0.000013518300038413145 s 0.00001392590001159988 s 0.97
value_and_grad / IPartOpt / cpu / Primal 0.000013710399907722604 s 0.000013735279972024728 s 1.00
value_and_grad / DefOpt / cpu / Primal 0.000014059000022825783 s 0.000014098959936745816 s 1.00
value_and_grad / IDefOpt / cpu / Primal 0.000013199100012570852 s 0.000014053879995117314 s 0.94
value_and_grad / JaXPipe / cuda / Primal 0.00003088 s 0.00003296 s 0.94
value_and_grad / Jax / cuda / Primal 0.000031104 s 0.000032800000000000004 s 0.95
value_and_grad / HLOOpt / cuda / Primal 0.000031808000000000004 s 0.000033472 s 0.95
value_and_grad / PartOpt / cuda / Primal 0.000031328 s 0.000033248 s 0.94
value_and_grad / IPartOpt / cuda / Primal 0.000031584 s 0.000033152000000000004 s 0.95
value_and_grad / DefOpt / cuda / Primal 0.000031584 s 0.000033408 s 0.95
value_and_grad / IDefOpt / cuda / Primal 0.000032032 s 0.000034176 s 0.94
value_and_grad / JaXPipe / tpu / Primal 0 s 0 s 1
value_and_grad / Jax / tpu / Primal 0 s 0 s 1
value_and_grad / HLOOpt / tpu / Primal 0 s 0 s 1
value_and_grad / PartOpt / tpu / Primal 0 s 0 s 1
value_and_grad / IPartOpt / tpu / Primal 0 s 0 s 1
value_and_grad / DefOpt / tpu / Primal 0 s 0 s 1
value_and_grad / IDefOpt / tpu / Primal 0 s 0 s 1
value_and_grad / JaXPipe / cpu / Primal 0.000023380000000000003 s 0.00001462693999201292 s 1.60
value_and_grad / Jax / cpu / Primal 0.00002254 s 0.000013887280019844183 s 1.62
value_and_grad / HLOOpt / cpu / Primal 0.000022735 s 0.000013906800022596144 s 1.63
value_and_grad / PartOpt / cpu / Primal 0.00002286 s 0.00001392590001159988 s 1.64
value_and_grad / IPartOpt / cpu / Primal 0.000022788 s 0.000013735279972024728 s 1.66
value_and_grad / DefOpt / cpu / Primal 0.00002281 s 0.000014098959936745816 s 1.62
value_and_grad / IDefOpt / cpu / Primal 0.000022567 s 0.000014053879995117314 s 1.61
jaxmd20 / JaXPipe / cuda / Primal 0.001457601 s 0.00146413 s 1.00
jaxmd20 / Jax / cuda / Primal 0.001494529 s 0.001500002 s 1.00
jaxmd20 / HLOOpt / cuda / Primal 0.001336768 s 0.001340257 s 1.00
jaxmd20 / PartOpt / cuda / Primal 0.00130656 s 0.001326465 s 0.98
jaxmd20 / IPartOpt / cuda / Primal 0.001358593 s 0.001347746 s 1.01
jaxmd20 / DefOpt / cuda / Primal 0.000928384 s 0.0009191679999999 s 1.01
jaxmd20 / IDefOpt / cuda / Primal 0.000947583 s 0.000950881 s 1.00
jaxmd20 / JaXPipe / cuda / Forward 0.001550976 s 0.001554785 s 1.00
jaxmd20 / Jax / cuda / Forward 0.002012225 s 0.0017800659999999 s 1.13
jaxmd20 / HLOOpt / cuda / Forward 0.001635072 s 0.001616514 s 1.01
jaxmd20 / PartOpt / cuda / Forward 0.001642048 s 0.001637218 s 1.00
jaxmd20 / IPartOpt / cuda / Forward 0.001623008 s 0.001613826 s 1.01
jaxmd20 / DefOpt / cuda / Forward 0.001651488 s 0.0016381149999999 s 1.01
jaxmd20 / IDefOpt / cuda / Forward 0.001623712 s 0.001622914 s 1.00
jaxmd20 / JaXPipe / cuda / PreRev 0.002686465 s 0.002663619 s 1.01
jaxmd20 / JaXPipe / cuda / PostRev 0.00529584 s 0.005329255 s 0.99
jaxmd20 / JaXPipe / cuda / BothRev 0.002713504 s 0.002686564 s 1.01
jaxmd20 / Jax / cuda / BothRev 0.005306689 s 0.005338087 s 0.99
jaxmd20 / HLOOpt / cuda / PreRev 0.002744736 s 0.002748932 s 1.00
jaxmd20 / HLOOpt / cuda / PostRev 0.005434882 s 0.005346055 s 1.02
jaxmd20 / HLOOpt / cuda / BothRev 0.002724864 s 0.002751748 s 0.99
jaxmd20 / PartOpt / cuda / PreRev 0.002863105 s 0.00281338 s 1.02
jaxmd20 / PartOpt / cuda / PostRev 0.005407777 s 0.0053920389999999 s 1.00
jaxmd20 / PartOpt / cuda / BothRev 0.0028226559999999 s 0.002791875 s 1.01
jaxmd20 / IPartOpt / cuda / PreRev 0.002836895 s 0.0028067549999999 s 1.01
jaxmd20 / IPartOpt / cuda / PostRev 0.005379586 s 0.005378982 s 1.00
jaxmd20 / IPartOpt / cuda / BothRev 0.002769857 s 0.002749955 s 1.01
jaxmd20 / DefOpt / cuda / PreRev 0.002836609 s 0.0028253149999999 s 1.00
jaxmd20 / DefOpt / cuda / PostRev 0.002745185 s 0.002760165 s 0.99
jaxmd20 / DefOpt / cuda / BothRev 0.002770209 s 0.0027866279999999 s 0.99
jaxmd20 / IDefOpt / cuda / PreRev 0.002831105 s 0.002807364 s 1.01
jaxmd20 / IDefOpt / cuda / PostRev 0.002322145 s 0.002303811 s 1.01
jaxmd20 / IDefOpt / cuda / BothRev 0.002766688 s 0.0027489629999999 s 1.01
jaxmd20 / JaXPipe / tpu / Primal 0.0092974662499999 s 0.0092963131249999 s 1.00
jaxmd20 / Jax / tpu / Primal 0.009269526875 s 0.00927079625 s 1.00
jaxmd20 / HLOOpt / tpu / Primal 0.009155528125 s 0.009157256875 s 1.00
jaxmd20 / PartOpt / tpu / Primal 0.009200965 s 0.00919751875 s 1.00
jaxmd20 / IPartOpt / tpu / Primal 0.009202034375 s 0.009203176875 s 1.00
jaxmd20 / DefOpt / tpu / Primal 0.0087966712499999 s 0.008796180625 s 1.00
jaxmd20 / IDefOpt / tpu / Primal 0.008704488125 s 0.008703285625 s 1.00
jaxmd20 / JaXPipe / tpu / Forward 0.017415119375 s 0.0174139475 s 1.00
jaxmd20 / Jax / tpu / Forward 0.01873873875 s 0.0187393075 s 1.00
jaxmd20 / HLOOpt / tpu / Forward 0.0174007875 s 0.0174003049999999 s 1.00
jaxmd20 / PartOpt / tpu / Forward 0.017416296875 s 0.017413809375 s 1.00
jaxmd20 / IPartOpt / tpu / Forward 0.01740998625 s 0.0174123799999999 s 1.00
jaxmd20 / DefOpt / tpu / Forward 0.017424361875 s 0.017426431875 s 1.00
jaxmd20 / IDefOpt / tpu / Forward 0.017411749375 s 0.017412701875 s 1.00
jaxmd20 / JaXPipe / tpu / PreRev 0.025447833125 s 0.0254493456249999 s 1.00
jaxmd20 / JaXPipe / tpu / PostRev 0.02187654375 s 0.021875376875 s 1.00
jaxmd20 / JaXPipe / tpu / BothRev 0.02547347375 s 0.025473811875 s 1.00
jaxmd20 / Jax / tpu / BothRev 0.0218740775 s 0.021873763125 s 1.00
jaxmd20 / HLOOpt / tpu / PreRev 0.025586626875 s 0.02558455 s 1.00
jaxmd20 / HLOOpt / tpu / PostRev 0.02071399125 s 0.020714709375 s 1.00
jaxmd20 / HLOOpt / tpu / BothRev 0.02569162875 s 0.025694736875 s 1.00
jaxmd20 / PartOpt / tpu / PreRev 0.025485230625 s 0.025487941875 s 1.00
jaxmd20 / PartOpt / tpu / PostRev 0.0212626956249999 s 0.02153581625 s 0.99
jaxmd20 / PartOpt / tpu / BothRev 0.02556918 s 0.0255683875 s 1.00
jaxmd20 / IPartOpt / tpu / PreRev 0.025474475625 s 0.0254793275 s 1.00
jaxmd20 / IPartOpt / tpu / PostRev 0.021519255 s 0.0215210425 s 1.00
jaxmd20 / IPartOpt / tpu / BothRev 0.02557093 s 0.02557064375 s 1.00
jaxmd20 / DefOpt / tpu / PreRev 0.02548422875 s 0.0254876425 s 1.00
jaxmd20 / DefOpt / tpu / PostRev 0.01882321125 s 0.018822664375 s 1.00
jaxmd20 / DefOpt / tpu / BothRev 0.025571416875 s 0.02557184625 s 1.00
jaxmd20 / IDefOpt / tpu / PreRev 0.025479161875 s 0.025477181875 s 1.00
jaxmd20 / IDefOpt / tpu / PostRev 0.0183412737499999 s 0.018344374375 s 1.00
jaxmd20 / IDefOpt / tpu / BothRev 0.025574326875 s 0.0255753725 s 1.00
jaxmd40 / JaXPipe / cpu / Primal 0.070437546 s 0.0755081639999999 s 0.93
jaxmd40 / Jax / cpu / Primal 0.069218235 s 0.0717200289999999 s 0.97
jaxmd40 / HLOOpt / cpu / Primal 0.0864528339999999 s 0.108675883 s 0.80
jaxmd40 / PartOpt / cpu / Primal 0.070029316 s 0.082305401 s 0.85
jaxmd40 / IPartOpt / cpu / Primal 0.064148335 s 0.084469558 s 0.76
jaxmd40 / DefOpt / cpu / Primal 0.086356319 s 0.10659383 s 0.81
jaxmd40 / IDefOpt / cpu / Primal 0.077289646 s 0.111599474 s 0.69
jaxmd40 / JaXPipe / cpu / Forward 0.160717668 s 0.189724921 s 0.85
jaxmd40 / Jax / cpu / Forward 0.083354955 s 0.089788619 s 0.93
jaxmd40 / HLOOpt / cpu / Forward 0.155012683 s 0.188932056 s 0.82
jaxmd40 / PartOpt / cpu / Forward 0.1641126899999999 s 0.187766515 s 0.87
jaxmd40 / IPartOpt / cpu / Forward 0.15625529 s 0.194974803 s 0.80
jaxmd40 / DefOpt / cpu / Forward 0.156870911 s 0.184983297 s 0.85
jaxmd40 / IDefOpt / cpu / Forward 0.152651269 s 0.190125706 s 0.80
jaxmd40 / JaXPipe / cpu / PreRev 0.211216208 s 0.251550023 s 0.84
jaxmd40 / JaXPipe / cpu / PostRev 0.139091416 s 0.155112812 s 0.90
jaxmd40 / JaXPipe / cpu / BothRev 0.229229774 s 0.248999614 s 0.92
jaxmd40 / Jax / cpu / BothRev 0.136818905 s 0.162482192 s 0.84
jaxmd40 / HLOOpt / cpu / PreRev 0.228181373 s 0.248515938 s 0.92
jaxmd40 / HLOOpt / cpu / PostRev 0.166616531 s 0.213123876 s 0.78
jaxmd40 / HLOOpt / cpu / BothRev 0.222540738 s 0.290243502 s 0.77
jaxmd40 / PartOpt / cpu / PreRev 0.20392956 s 0.243762015 s 0.84
jaxmd40 / PartOpt / cpu / PostRev 0.128692598 s 0.145630436 s 0.88
jaxmd40 / PartOpt / cpu / BothRev 0.224388404 s 0.267546272 s 0.84
jaxmd40 / IPartOpt / cpu / PreRev 0.2099655899999999 s 0.246754119 s 0.85
jaxmd40 / IPartOpt / cpu / PostRev 0.135609714 s 0.134807949 s 1.01
jaxmd40 / IPartOpt / cpu / BothRev 0.25953237 s 0.2652395629999999 s 0.98
jaxmd40 / DefOpt / cpu / PreRev 0.216519988 s 0.255335409 s 0.85
jaxmd40 / DefOpt / cpu / PostRev 0.17726939 s 0.216085604 s 0.82
jaxmd40 / DefOpt / cpu / BothRev 0.2416942819999999 s 0.287981373 s 0.84
jaxmd40 / IDefOpt / cpu / PreRev 0.233331476 s 0.249678795 s 0.93
jaxmd40 / IDefOpt / cpu / PostRev 0.1675968159999999 s 0.210582048 s 0.80
jaxmd40 / IDefOpt / cpu / BothRev 0.250686033 s 0.26354349 s 0.95
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / cuda / Primal 1.701665501 s 1.703403109 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / cuda / Primal 1.703885154 s 1.704959474 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / cuda / Primal 1.715888839 s 1.714279019 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / cuda / Primal 1.696104386 s 1.695521465 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / cuda / Primal 1.694128773 s 1.694030017 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / cuda / Primal 1.665237092 s 1.664191437 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / cuda / Primal 1.912735816 s 1.914000938 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / tpu / Primal 3.038123845625 s 3.038906811875 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / tpu / Primal 3.038594565 s 3.03936133 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / tpu / Primal 3.120960913125 s 3.12152741125 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / tpu / Primal 3.05947385125 s 3.060072164375 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / tpu / Primal 3.059731680625 s 3.060387236875 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / tpu / Primal 2.102142330625 s 2.10243730625 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / tpu / Primal 2.94743460375 s 2.94835662125 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / JaXPipe / cpu / Primal 5.955080135 s 6.770451712 s 0.88
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / Jax / cpu / Primal 5.868964441 s 6.796283552 s 0.86
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / HLOOpt / cpu / Primal 5.913522663 s 6.784701461 s 0.87
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / PartOpt / cpu / Primal 6.024176476 s 6.773146511 s 0.89
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IPartOpt / cpu / Primal 5.972733144 s 6.831082725 s 0.87
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / DefOpt / cpu / Primal 2.312797727 s 2.680127355 s 0.86
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IDefOpt / cpu / Primal 6.456180052 s 7.33637517 s 0.88

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

@gaurav-arya gaurav-arya force-pushed the ag/ndim_symmetry_lattice branch from c59e2bb to d164c18 Compare December 1, 2025 21:08
@gaurav-arya gaurav-arya force-pushed the ag/ndim_symmetry_lattice branch from 53945b8 to e62ad33 Compare December 1, 2025 21:47
avik-pal
avik-pal reviewed Dec 26, 2025
View reviewed changes
Copy link
Collaborator

@avik-pal avik-pal left a comment

Choose a reason for hiding this comment

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

Generally LGTM. For future, I would add a small section above the partial symmetry annotation describing what the attribute means mathematically

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@avik-pal avik-pal avik-pal left review comments

@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.

4 participants

@gaurav-arya @wsmoses @avik-pal