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 24 commits into
base: main
Choose a base branch
from
Open

Add partial symmetry detection #1663

gaurav-arya wants to merge 24 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: 93f3090 Previous: 66133db Ratio
actmtch / JaXPipe / cpu / Primal 0.000007168500060288352 s 0.00000713281997377635 s 1.01
actmtch / Jax / cpu / Primal 0.000006658040028924006 s 0.000006339999999909196 s 1.05
actmtch / HLOOpt / cpu / Primal 0.000008639440038678003 s 0.000007280080017153523 s 1.19
actmtch / PartOpt / cpu / Primal 0.000007306679999601329 s 0.000006684480013063876 s 1.09
actmtch / IPartOpt / cpu / Primal 0.000007142780032154406 s 0.000006549319987243508 s 1.09
actmtch / DefOpt / cpu / Primal 0.000007801479969202773 s 0.000007395840020762989 s 1.05
actmtch / IDefOpt / cpu / Primal 0.000007775099948048592 s 0.000007337020006161765 s 1.06
actmtch / JaXPipe / cpu / Forward 0.000012040219980917754 s 0.00001092381994567404 s 1.10
actmtch / Jax / cpu / Forward 0.000010655880005288057 s 0.000009964540013243094 s 1.07
actmtch / HLOOpt / cpu / Forward 0.000012388940012897364 s 0.000010882579999815787 s 1.14
actmtch / PartOpt / cpu / Forward 0.000012165760026618954 s 0.000010533580052651814 s 1.15
actmtch / IPartOpt / cpu / Forward 0.000012523580026027047 s 0.000010967720008920878 s 1.14
actmtch / DefOpt / cpu / Forward 0.000011571519999051817 s 0.00001072210001439089 s 1.08
actmtch / IDefOpt / cpu / Forward 0.000011615619969234104 s 0.00001106340000660566 s 1.05
actmtch / JaXPipe / cpu / PreRev 0.000011783960044340348 s 0.000011070179989474126 s 1.06
actmtch / JaXPipe / cpu / PostRev 0.000011226519982301395 s 0.000009969440015993314 s 1.13
actmtch / JaXPipe / cpu / BothRev 0.000012999620030313964 s 0.00001166578003903851 s 1.11
actmtch / Jax / cpu / BothRev 0.000010355979975429364 s 0.0000100311000005604 s 1.03
actmtch / HLOOpt / cpu / PreRev 0.00001168239998150966 s 0.000011173419934493725 s 1.05
actmtch / HLOOpt / cpu / PostRev 0.000014169260030030271 s 0.00001262077998035238 s 1.12
actmtch / HLOOpt / cpu / BothRev 0.000013062359985269722 s 0.000010960199988403474 s 1.19
actmtch / PartOpt / cpu / PreRev 0.00001152832003754156 s 0.000010591719992589788 s 1.09
actmtch / PartOpt / cpu / PostRev 0.000011118079974039577 s 0.000010557359992162674 s 1.05
actmtch / PartOpt / cpu / BothRev 0.00001278959999581275 s 0.00001126127997849835 s 1.14
actmtch / IPartOpt / cpu / PreRev 0.000011939739961235318 s 0.000010461360016051911 s 1.14
actmtch / IPartOpt / cpu / PostRev 0.000011389540040909195 s 0.000009731580012157791 s 1.17
actmtch / IPartOpt / cpu / BothRev 0.000012631839999812656 s 0.000010840979966815211 s 1.17
actmtch / DefOpt / cpu / PreRev 0.000011740379995899275 s 0.00001129805999880773 s 1.04
actmtch / DefOpt / cpu / PostRev 0.000012289319984120085 s 0.000010375859956184286 s 1.18
actmtch / DefOpt / cpu / BothRev 0.000013230639988250914 s 0.000010885640003834853 s 1.22
actmtch / IDefOpt / cpu / PreRev 0.000011605799991230014 s 0.000010767719995783407 s 1.08
actmtch / IDefOpt / cpu / PostRev 0.00001278891998481413 s 0.000011479120003059506 s 1.11
actmtch / IDefOpt / cpu / BothRev 0.0000125474799915537 s 0.000010819939989232807 s 1.16
actmtch / JaXPipe / cuda / Primal 0.000002431 s 0.000002015 s 1.21
actmtch / Jax / cuda / Primal 0.000002431 s 0.000002015 s 1.21
actmtch / HLOOpt / cuda / Primal 0.0000024 s 0.000002015 s 1.19
actmtch / PartOpt / cuda / Primal 0.0000024 s 0.000002015 s 1.19
actmtch / IPartOpt / cuda / Primal 0.0000024 s 0.000002015 s 1.19
actmtch / DefOpt / cuda / Primal 0.000002431 s 0.000002015 s 1.21
actmtch / IDefOpt / cuda / Primal 0.000002431 s 0.000002015 s 1.21
actmtch / JaXPipe / cuda / Forward 0.000011455999999999998 s 0.000009856 s 1.16
actmtch / Jax / cuda / Forward 0.0000104 s 0.000009856 s 1.06
actmtch / HLOOpt / cuda / Forward 0.000010432 s 0.000009696 s 1.08
actmtch / PartOpt / cuda / Forward 0.000010528 s 0.000009952 s 1.06
actmtch / IPartOpt / cuda / Forward 0.000010624 s 0.000009824 s 1.08
actmtch / DefOpt / cuda / Forward 0.000010176 s 0.000010112 s 1.01
actmtch / IDefOpt / cuda / Forward 0.000010496 s 0.00001008 s 1.04
actmtch / JaXPipe / cuda / PreRev 0.000010432 s 0.00000992 s 1.05
actmtch / JaXPipe / cuda / PostRev 0.00001072 s 0.000009985 s 1.07
actmtch / JaXPipe / cuda / BothRev 0.000010688 s 0.000010144 s 1.05
actmtch / Jax / cuda / BothRev 0.000010623 s 0.000010272 s 1.03
actmtch / HLOOpt / cuda / PreRev 0.000010592 s 0.000010432 s 1.02
actmtch / HLOOpt / cuda / PostRev 0.000010591 s 0.000010016 s 1.06
actmtch / HLOOpt / cuda / BothRev 0.000010656 s 0.000010368 s 1.03
actmtch / PartOpt / cuda / PreRev 0.00001072 s 0.000009888 s 1.08
actmtch / PartOpt / cuda / PostRev 0.000010336 s 0.00001008 s 1.03
actmtch / PartOpt / cuda / BothRev 0.000010336 s 0.000010209 s 1.01
actmtch / IPartOpt / cuda / PreRev 0.000010464 s 0.000010079 s 1.04
actmtch / IPartOpt / cuda / PostRev 0.000011296 s 0.000010208 s 1.11
actmtch / IPartOpt / cuda / BothRev 0.000010048 s 0.00000992 s 1.01
actmtch / DefOpt / cuda / PreRev 0.000010624 s 0.000010303 s 1.03
actmtch / DefOpt / cuda / PostRev 0.000010432 s 0.000009888 s 1.06
actmtch / DefOpt / cuda / BothRev 0.000010592 s 0.000009953 s 1.06
actmtch / IDefOpt / cuda / PreRev 0.000010592 s 0.00001056 s 1.00
actmtch / IDefOpt / cuda / PostRev 0.000011551 s 0.000010272 s 1.12
actmtch / IDefOpt / cuda / BothRev 0.000010751 s 0.000010336 s 1.04
actmtch / JaXPipe / tpu / Primal 5.6315e-7 s 5.632000000000001e-7 s 1.00
actmtch / Jax / tpu / Primal 5.965499999999999e-7 s 5.97e-7 s 1.00
actmtch / HLOOpt / tpu / Primal 0.000002092725 s 0.000002100575 s 1.00
actmtch / PartOpt / tpu / Primal 5.96575e-7 s 5.964250000000001e-7 s 1.00
actmtch / IPartOpt / tpu / Primal 5.527750000000001e-7 s 5.525e-7 s 1.00
actmtch / DefOpt / tpu / Primal 0.00000217245 s 0.0000021614 s 1.01
actmtch / IDefOpt / tpu / Primal 0.000002101025 s 0.000002094925 s 1.00
actmtch / JaXPipe / tpu / Forward 0.000003827325 s 0.000003830450000000001 s 1.00
actmtch / Jax / tpu / Forward 0.0000012096750000000002 s 0.000001206 s 1.00
actmtch / HLOOpt / tpu / Forward 0.00000394575 s 0.0000039277000000000005 s 1.00
actmtch / PartOpt / tpu / Forward 0.000003916025 s 0.0000039145500000000005 s 1.00
actmtch / IPartOpt / tpu / Forward 0.000003943725 s 0.000003941225 s 1.00
actmtch / DefOpt / tpu / Forward 0.000003918925 s 0.0000039124 s 1.00
actmtch / IDefOpt / tpu / Forward 0.000003933749999999999 s 0.0000039432 s 1.00
actmtch / JaXPipe / tpu / PreRev 0.000003482225 s 0.0000034772 s 1.00
actmtch / JaXPipe / tpu / PostRev 0.00000163675 s 0.000001644325 s 1.00
actmtch / JaXPipe / tpu / BothRev 0.00000348215 s 0.00000347105 s 1.00
actmtch / Jax / tpu / BothRev 0.0000016368250000000002 s 0.0000016378999999999998 s 1.00
actmtch / HLOOpt / tpu / PreRev 0.0000034929 s 0.000003480525 s 1.00
actmtch / HLOOpt / tpu / PostRev 0.000003400825 s 0.0000034096 s 1.00
actmtch / HLOOpt / tpu / BothRev 0.0000034759 s 0.00000346665 s 1.00
actmtch / PartOpt / tpu / PreRev 0.0000034156750000000005 s 0.0000034179 s 1.00
actmtch / PartOpt / tpu / PostRev 0.00000159895 s 0.0000015853 s 1.01
actmtch / PartOpt / tpu / BothRev 0.00000340945 s 0.000003403625 s 1.00
actmtch / IPartOpt / tpu / PreRev 0.0000034935 s 0.0000034651749999999994 s 1.01
actmtch / IPartOpt / tpu / PostRev 0.0000016436999999999998 s 0.00000163315 s 1.01
actmtch / IPartOpt / tpu / BothRev 0.0000034882249999999995 s 0.000003482425 s 1.00
actmtch / DefOpt / tpu / PreRev 0.0000034037 s 0.000003414375 s 1.00
actmtch / DefOpt / tpu / PostRev 0.000003419 s 0.0000034209750000000003 s 1.00
actmtch / DefOpt / tpu / BothRev 0.000003404875 s 0.0000034067750000000005 s 1.00
actmtch / IDefOpt / tpu / PreRev 0.00000348545 s 0.00000347925 s 1.00
actmtch / IDefOpt / tpu / PostRev 0.00000342395 s 0.000003409275 s 1.00
actmtch / IDefOpt / tpu / BothRev 0.0000034821 s 0.0000034686 s 1.00
actmtch / JaXPipe / cpu / Primal 0.000013086 s 0.00000713281997377635 s 1.83
actmtch / Jax / cpu / Primal 0.000013117 s 0.000006339999999909196 s 2.07
actmtch / HLOOpt / cpu / Primal 0.000013863 s 0.000007280080017153523 s 1.90
actmtch / PartOpt / cpu / Primal 0.00001311 s 0.000006684480013063876 s 1.96
actmtch / IPartOpt / cpu / Primal 0.000013057 s 0.000006549319987243508 s 1.99
actmtch / DefOpt / cpu / Primal 0.000013826 s 0.000007395840020762989 s 1.87
actmtch / IDefOpt / cpu / Primal 0.000013948 s 0.000007337020006161765 s 1.90
actmtch / JaXPipe / cpu / Forward 0.000019516 s 0.00001092381994567404 s 1.79
actmtch / Jax / cpu / Forward 0.000018361 s 0.000009964540013243094 s 1.84
actmtch / HLOOpt / cpu / Forward 0.000019315 s 0.000010882579999815787 s 1.77
actmtch / PartOpt / cpu / Forward 0.000019211 s 0.000010533580052651814 s 1.82
actmtch / IPartOpt / cpu / Forward 0.000019187 s 0.000010967720008920878 s 1.75
actmtch / DefOpt / cpu / Forward 0.000019023 s 0.00001072210001439089 s 1.77
actmtch / IDefOpt / cpu / Forward 0.000019435 s 0.00001106340000660566 s 1.76
actmtch / JaXPipe / cpu / PreRev 0.000019901 s 0.000011070179989474126 s 1.80
actmtch / JaXPipe / cpu / PostRev 0.00001772 s 0.000009969440015993314 s 1.78
actmtch / JaXPipe / cpu / BothRev 0.000019525 s 0.00001166578003903851 s 1.67
actmtch / Jax / cpu / BothRev 0.000018304 s 0.0000100311000005604 s 1.82
actmtch / HLOOpt / cpu / PreRev 0.000018871 s 0.000011173419934493725 s 1.69
actmtch / HLOOpt / cpu / PostRev 0.000019695 s 0.00001262077998035238 s 1.56
actmtch / HLOOpt / cpu / BothRev 0.000019429 s 0.000010960199988403474 s 1.77
actmtch / PartOpt / cpu / PreRev 0.000019333 s 0.000010591719992589788 s 1.83
actmtch / PartOpt / cpu / PostRev 0.000017628 s 0.000010557359992162674 s 1.67
actmtch / PartOpt / cpu / BothRev 0.000019821 s 0.00001126127997849835 s 1.76
actmtch / IPartOpt / cpu / PreRev 0.000019218 s 0.000010461360016051911 s 1.84
actmtch / IPartOpt / cpu / PostRev 0.00001789 s 0.000009731580012157791 s 1.84
actmtch / IPartOpt / cpu / BothRev 0.000018622 s 0.000010840979966815211 s 1.72
actmtch / DefOpt / cpu / PreRev 0.000019408 s 0.00001129805999880773 s 1.72
actmtch / DefOpt / cpu / PostRev 0.000019538 s 0.000010375859956184286 s 1.88
actmtch / DefOpt / cpu / BothRev 0.000019503 s 0.000010885640003834853 s 1.79
actmtch / IDefOpt / cpu / PreRev 0.00001934 s 0.000010767719995783407 s 1.80
actmtch / IDefOpt / cpu / PostRev 0.000019348 s 0.000011479120003059506 s 1.69
actmtch / IDefOpt / cpu / BothRev 0.000020063 s 0.000010819939989232807 s 1.85
add_one / JaXPipe / cpu / Primal 0.000007465920025424566 s 0.000006659000000581727 s 1.12
add_one / Jax / cpu / Primal 0.000007198139974207151 s 0.000006576479972864036 s 1.09
add_one / HLOOpt / cpu / Primal 0.00000776547997702437 s 0.0000066848000187746944 s 1.16
add_one / PartOpt / cpu / Primal 0.000007272339944393025 s 0.000006537359977301093 s 1.11
add_one / IPartOpt / cpu / Primal 0.000007727280035396689 s 0.000006783579992770683 s 1.14
add_one / DefOpt / cpu / Primal 0.000006944319984540926 s 0.000006492540005638148 s 1.07
add_one / IDefOpt / cpu / Primal 0.000007189320049292292 s 0.000006682360044578673 s 1.08
add_one / JaXPipe / cpu / Forward 0.00001069335997271992 s 0.000009964399996533756 s 1.07
add_one / Jax / cpu / Forward 0.000010391739997430704 s 0.000010458340038894676 s 0.99
add_one / HLOOpt / cpu / Forward 0.00001084305999938806 s 0.000010503659987080028 s 1.03
add_one / PartOpt / cpu / Forward 0.000010767320009108515 s 0.000009928980007316569 s 1.08
add_one / IPartOpt / cpu / Forward 0.00001075270001820172 s 0.0000104195800031448 s 1.03
add_one / DefOpt / cpu / Forward 0.000010515319982005168 s 0.000010610659956000744 s 0.99
add_one / IDefOpt / cpu / Forward 0.000011358800002199132 s 0.000010212440029135906 s 1.11
add_one / JaXPipe / cpu / PreRev 0.000012953839959664038 s 0.000011650119986370556 s 1.11
add_one / JaXPipe / cpu / PostRev 0.000012641759949474365 s 0.00001176739995571552 s 1.07
add_one / JaXPipe / cpu / BothRev 0.00001343293994978012 s 0.000012446319997252431 s 1.08
add_one / Jax / cpu / BothRev 0.000013159979998818016 s 0.000011605139998209778 s 1.13
add_one / HLOOpt / cpu / PreRev 0.000013052519989287248 s 0.00001167976000942872 s 1.12
add_one / HLOOpt / cpu / PostRev 0.000014555739971910952 s 0.000013972540000395385 s 1.04
add_one / HLOOpt / cpu / BothRev 0.00001235493997228332 s 0.000011715100017681834 s 1.05
add_one / PartOpt / cpu / PreRev 0.00001266584001314186 s 0.000011190120021637997 s 1.13
add_one / PartOpt / cpu / PostRev 0.000013158179990568895 s 0.000011890860023413551 s 1.11
add_one / PartOpt / cpu / BothRev 0.00001405035997777304 s 0.00001255587994819507 s 1.12
add_one / IPartOpt / cpu / PreRev 0.000012919300033900071 s 0.000011687179994623876 s 1.11
add_one / IPartOpt / cpu / PostRev 0.000012487800004237216 s 0.000011444279998613638 s 1.09
add_one / IPartOpt / cpu / BothRev 0.000012756100031765528 s 0.000011512159999256254 s 1.11
add_one / DefOpt / cpu / PreRev 0.000012389880021146384 s 0.000011918979998881695 s 1.04
add_one / DefOpt / cpu / PostRev 0.000012959239966221502 s 0.000011530939973454224 s 1.12
add_one / DefOpt / cpu / BothRev 0.000012935139966430143 s 0.000011958420000155456 s 1.08
add_one / IDefOpt / cpu / PreRev 0.000013050919960733154 s 0.000011380919968360104 s 1.15
add_one / IDefOpt / cpu / PostRev 0.00001249692003511882 s 0.00001230832002875104 s 1.02
add_one / IDefOpt / cpu / BothRev 0.000012610899984792924 s 0.000011778740017689416 s 1.07
add_one / JaXPipe / cuda / Primal 0.000002335 s 0.000001919 s 1.22
add_one / Jax / cuda / Primal 0.000002335 s 0.0000019200000000000003 s 1.22
add_one / HLOOpt / cuda / Primal 0.000002335 s 0.0000019200000000000003 s 1.22
add_one / PartOpt / cuda / Primal 0.000002335 s 0.0000019200000000000003 s 1.22
add_one / IPartOpt / cuda / Primal 0.000002335 s 0.0000019200000000000003 s 1.22
add_one / DefOpt / cuda / Primal 0.000002335 s 0.0000019200000000000003 s 1.22
add_one / IDefOpt / cuda / Primal 0.000002335 s 0.0000019200000000000003 s 1.22
add_one / JaXPipe / cuda / Forward 0.000010496 s 0.00001056 s 0.99
add_one / Jax / cuda / Forward 0.000010688 s 0.000010144 s 1.05
add_one / HLOOpt / cuda / Forward 0.000010048 s 0.000009888 s 1.02
add_one / PartOpt / cuda / Forward 0.000010625 s 0.000009952 s 1.07
add_one / IPartOpt / cuda / Forward 0.000010432 s 0.000010016 s 1.04
add_one / DefOpt / cuda / Forward 0.000010624 s 0.000010144 s 1.05
add_one / IDefOpt / cuda / Forward 0.000010369 s 0.000010175 s 1.02
add_one / JaXPipe / cuda / PreRev 0.000025887 s 0.000025216 s 1.03
add_one / JaXPipe / cuda / PostRev 0.000025472000000000003 s 0.000024896 s 1.02
add_one / JaXPipe / cuda / BothRev 0.00002576 s 0.000025088 s 1.03
add_one / Jax / cuda / BothRev 0.000025312 s 0.00002608 s 0.97
add_one / HLOOpt / cuda / PreRev 0.000026143 s 0.000025184 s 1.04
add_one / HLOOpt / cuda / PostRev 0.000025696 s 0.00002464 s 1.04
add_one / HLOOpt / cuda / BothRev 0.000025472000000000003 s 0.00002464 s 1.03
add_one / PartOpt / cuda / PreRev 0.000025728 s 0.000024992 s 1.03
add_one / PartOpt / cuda / PostRev 0.000025471 s 0.000025472000000000003 s 1.00
add_one / PartOpt / cuda / BothRev 0.000025439 s 0.00002528 s 1.01
add_one / IPartOpt / cuda / PreRev 0.000025823 s 0.000025023 s 1.03
add_one / IPartOpt / cuda / PostRev 0.000025887 s 0.00002512 s 1.03
add_one / IPartOpt / cuda / BothRev 0.000026016 s 0.000025408 s 1.02
add_one / DefOpt / cuda / PreRev 0.000025823 s 0.000025216 s 1.02
add_one / DefOpt / cuda / PostRev 0.00002576 s 0.000024928 s 1.03
add_one / DefOpt / cuda / BothRev 0.000025983 s 0.000025216 s 1.03
add_one / IDefOpt / cuda / PreRev 0.000025727 s 0.000025248 s 1.02
add_one / IDefOpt / cuda / PostRev 0.000026112 s 0.000025408 s 1.03
add_one / IDefOpt / cuda / BothRev 0.000025952 s 0.000024992 s 1.04
add_one / JaXPipe / tpu / Primal 0.0000013990749999999998 s 0.0000014259250000000002 s 0.98
add_one / Jax / tpu / Primal 0.0000014021749999999998 s 0.0000014046 s 1.00
add_one / HLOOpt / tpu / Primal 0.0000014115249999999998 s 0.0000014235999999999998 s 0.99
add_one / PartOpt / tpu / Primal 0.0000013958 s 0.0000014038999999999998 s 0.99
add_one / IPartOpt / tpu / Primal 0.0000014040749999999998 s 0.00000142825 s 0.98
add_one / DefOpt / tpu / Primal 0.000001394175 s 0.0000014073749999999998 s 0.99
add_one / IDefOpt / tpu / Primal 0.000001401575 s 0.0000014244 s 0.98
add_one / JaXPipe / tpu / Forward 0.0000017914 s 0.00000185515 s 0.97
add_one / Jax / tpu / Forward 0.0000018189 s 0.0000018396 s 0.99
add_one / HLOOpt / tpu / Forward 0.0000017971500000000002 s 0.0000018493 s 0.97
add_one / PartOpt / tpu / Forward 0.00000181705 s 0.00000183585 s 0.99
add_one / IPartOpt / tpu / Forward 0.000001794575 s 0.000001853275 s 0.97
add_one / DefOpt / tpu / Forward 0.000001817 s 0.0000018479 s 0.98
add_one / IDefOpt / tpu / Forward 0.0000017960999999999998 s 0.000001846525 s 0.97
add_one / JaXPipe / tpu / PreRev 0.0000022207 s 0.0000022328 s 0.99
add_one / JaXPipe / tpu / PostRev 0.000002183025 s 0.0000022363750000000003 s 0.98
add_one / JaXPipe / tpu / BothRev 0.0000022107250000000003 s 0.000002242525 s 0.99
add_one / Jax / tpu / BothRev 0.0000021815 s 0.0000022507 s 0.97
add_one / HLOOpt / tpu / PreRev 0.00000221995 s 0.00000225035 s 0.99
add_one / HLOOpt / tpu / PostRev 0.000002184475 s 0.0000022378 s 0.98
add_one / HLOOpt / tpu / BothRev 0.0000022118750000000003 s 0.0000022373 s 0.99
add_one / PartOpt / tpu / PreRev 0.000002179325 s 0.000002240775 s 0.97
add_one / PartOpt / tpu / PostRev 0.0000022151750000000004 s 0.0000022342500000000003 s 0.99
add_one / PartOpt / tpu / BothRev 0.00000218105 s 0.000002242075 s 0.97
add_one / IPartOpt / tpu / PreRev 0.000002219475 s 0.0000022358500000000004 s 0.99
add_one / IPartOpt / tpu / PostRev 0.000002188175 s 0.000002243625 s 0.98
add_one / IPartOpt / tpu / BothRev 0.000002211975 s 0.0000022312 s 0.99
add_one / DefOpt / tpu / PreRev 0.00000219035 s 0.0000022413 s 0.98
add_one / DefOpt / tpu / PostRev 0.0000022151750000000004 s 0.0000022313 s 0.99
add_one / DefOpt / tpu / BothRev 0.000002196025 s 0.00000223185 s 0.98
add_one / IDefOpt / tpu / PreRev 0.000002230875 s 0.00000224095 s 1.00
add_one / IDefOpt / tpu / PostRev 0.000002191075 s 0.00000224465 s 0.98
add_one / IDefOpt / tpu / BothRev 0.000002210525 s 0.000002237925 s 0.99
add_one / JaXPipe / cpu / Primal 0.000013285 s 0.000006659000000581727 s 2.00
add_one / Jax / cpu / Primal 0.000012796 s 0.000006576479972864036 s 1.95
add_one / HLOOpt / cpu / Primal 0.000013344999999999998 s 0.0000066848000187746944 s 2.00
add_one / PartOpt / cpu / Primal 0.000012925 s 0.000006537359977301093 s 1.98
add_one / IPartOpt / cpu / Primal 0.000012787 s 0.000006783579992770683 s 1.88
add_one / DefOpt / cpu / Primal 0.000012837 s 0.000006492540005638148 s 1.98
add_one / IDefOpt / cpu / Primal 0.000012651 s 0.000006682360044578673 s 1.89
add_one / JaXPipe / cpu / Forward 0.000018011 s 0.000009964399996533756 s 1.81
add_one / Jax / cpu / Forward 0.000017551000000000002 s 0.000010458340038894676 s 1.68
add_one / HLOOpt / cpu / Forward 0.000017743 s 0.000010503659987080028 s 1.69
add_one / PartOpt / cpu / Forward 0.000017902000000000002 s 0.000009928980007316569 s 1.80
add_one / IPartOpt / cpu / Forward 0.00001795 s 0.0000104195800031448 s 1.72
add_one / DefOpt / cpu / Forward 0.000017718999999999998 s 0.000010610659956000744 s 1.67
add_one / IDefOpt / cpu / Forward 0.000017854 s 0.000010212440029135906 s 1.75
add_one / JaXPipe / cpu / PreRev 0.000019869 s 0.000011650119986370556 s 1.71
add_one / JaXPipe / cpu / PostRev 0.000019694 s 0.00001176739995571552 s 1.67
add_one / JaXPipe / cpu / BothRev 0.000019914000000000003 s 0.000012446319997252431 s 1.60
add_one / Jax / cpu / BothRev 0.000019215 s 0.000011605139998209778 s 1.66
add_one / HLOOpt / cpu / PreRev 0.000019653 s 0.00001167976000942872 s 1.68
add_one / HLOOpt / cpu / PostRev 0.000020059 s 0.000013972540000395385 s 1.44
add_one / HLOOpt / cpu / BothRev 0.00001984 s 0.000011715100017681834 s 1.69
add_one / PartOpt / cpu / PreRev 0.000019678 s 0.000011190120021637997 s 1.76
add_one / PartOpt / cpu / PostRev 0.000019987 s 0.000011890860023413551 s 1.68
add_one / PartOpt / cpu / BothRev 0.000019367 s 0.00001255587994819507 s 1.54
add_one / IPartOpt / cpu / PreRev 0.000019851 s 0.000011687179994623876 s 1.70
add_one / IPartOpt / cpu / PostRev 0.000019331 s 0.000011444279998613638 s 1.69
add_one / IPartOpt / cpu / BothRev 0.000019898 s 0.000011512159999256254 s 1.73
add_one / DefOpt / cpu / PreRev 0.000019823 s 0.000011918979998881695 s 1.66
add_one / DefOpt / cpu / PostRev 0.000020165 s 0.000011530939973454224 s 1.75
add_one / DefOpt / cpu / BothRev 0.000019511 s 0.000011958420000155456 s 1.63
add_one / IDefOpt / cpu / PreRev 0.000019582 s 0.000011380919968360104 s 1.72
add_one / IDefOpt / cpu / PostRev 0.000019843 s 0.00001230832002875104 s 1.61
add_one / IDefOpt / cpu / BothRev 0.00001948 s 0.000011778740017689416 s 1.65
add_two / JaXPipe / cpu / Primal 0.0000069443800202861895 s 0.000006892639967190917 s 1.01
add_two / Jax / cpu / Primal 0.00000770147998991888 s 0.000006933860004210146 s 1.11
add_two / HLOOpt / cpu / Primal 0.000007871180041547632 s 0.000007509379993280163 s 1.05
add_two / PartOpt / cpu / Primal 0.0000070648600194545 s 0.000006904719966769335 s 1.02
add_two / IPartOpt / cpu / Primal 0.00000732825996237807 s 0.000007135020032364992 s 1.03
add_two / DefOpt / cpu / Primal 0.0000069412399898283184 s 0.000006958459998713806 s 1.00
add_two / IDefOpt / cpu / Primal 0.000007213999988380238 s 0.0000065641800119919935 s 1.10
add_two / JaXPipe / cpu / Forward 0.000011463740002000122 s 0.000010514580008020858 s 1.09
add_two / Jax / cpu / Forward 0.000011033580021830855 s 0.000010020859963333351 s 1.10
add_two / HLOOpt / cpu / Forward 0.00001119774000471807 s 0.000010757099989859853 s 1.04
add_two / PartOpt / cpu / Forward 0.00001098245998036873 s 0.000010200499991697142 s 1.08
add_two / IPartOpt / cpu / Forward 0.000010946459970000432 s 0.000010306379981557256 s 1.06
add_two / DefOpt / cpu / Forward 0.000011172720005561132 s 0.00001038984000842902 s 1.08
add_two / IDefOpt / cpu / Forward 0.00001091814002393221 s 0.000010215420024906052 s 1.07
add_two / JaXPipe / cpu / PreRev 0.000015474219981115312 s 0.000013841239997418598 s 1.12
add_two / JaXPipe / cpu / PostRev 0.000014991040025051916 s 0.000014024039992364124 s 1.07
add_two / JaXPipe / cpu / BothRev 0.00001527571998849453 s 0.000014294040029199096 s 1.07
add_two / Jax / cpu / BothRev 0.000015382199981104348 s 0.000013598579989775315 s 1.13
add_two / HLOOpt / cpu / PreRev 0.000015323160032494342 s 0.00001361135996376106 s 1.13
add_two / HLOOpt / cpu / PostRev 0.000016962019972197596 s 0.000016257539991784144 s 1.04
add_two / HLOOpt / cpu / BothRev 0.00001573665998876095 s 0.000014164219992380822 s 1.11
add_two / PartOpt / cpu / PreRev 0.000015736499972263118 s 0.000013886520009691596 s 1.13
add_two / PartOpt / cpu / PostRev 0.000015602879975631368 s 0.00001393359996654908 s 1.12
add_two / PartOpt / cpu / BothRev 0.000015640979991076166 s 0.00001496692000728217 s 1.05
add_two / IPartOpt / cpu / PreRev 0.000015566800011583838 s 0.000013847800000803544 s 1.12
add_two / IPartOpt / cpu / PostRev 0.000014657100018666825 s 0.000014076920015213546 s 1.04
add_two / IPartOpt / cpu / BothRev 0.00001526023997939774 s 0.000014492899990727893 s 1.05
add_two / DefOpt / cpu / PreRev 0.000015033500048957648 s 0.00001408632003403909 s 1.07
add_two / DefOpt / cpu / PostRev 0.00001571133995639684 s 0.000014134579969322658 s 1.11
add_two / DefOpt / cpu / BothRev 0.000016293979988404318 s 0.000014213100002962163 s 1.15
add_two / IDefOpt / cpu / PreRev 0.000015565579997200984 s 0.000013673720004589996 s 1.14
add_two / IDefOpt / cpu / PostRev 0.000015456459959750645 s 0.000014411659985853476 s 1.07
add_two / IDefOpt / cpu / BothRev 0.000016703699984645936 s 0.0000148106800406822 s 1.13
add_two / JaXPipe / cuda / Primal 0.000002432 s 0.000001887 s 1.29
add_two / Jax / cuda / Primal 0.000002432 s 0.000001887 s 1.29
add_two / HLOOpt / cuda / Primal 0.000002431 s 0.000001887 s 1.29
add_two / PartOpt / cuda / Primal 0.000002432 s 0.000001888 s 1.29
add_two / IPartOpt / cuda / Primal 0.000002432 s 0.000001887 s 1.29
add_two / DefOpt / cuda / Primal 0.000002432 s 0.000001887 s 1.29
add_two / IDefOpt / cuda / Primal 0.000002432 s 0.000001887 s 1.29
add_two / JaXPipe / cuda / Forward 0.000010496 s 0.000010048 s 1.04
add_two / Jax / cuda / Forward 0.000010592 s 0.000009824 s 1.08
add_two / HLOOpt / cuda / Forward 0.000010592 s 0.0000096 s 1.10
add_two / PartOpt / cuda / Forward 0.0000104 s 0.000009696 s 1.07
add_two / IPartOpt / cuda / Forward 0.000010112 s 0.000009985 s 1.01
add_two / DefOpt / cuda / Forward 0.000010656 s 0.000009792 s 1.09
add_two / IDefOpt / cuda / Forward 0.000010496 s 0.000009727 s 1.08
add_two / JaXPipe / cuda / PreRev 0.000033568 s 0.000031519 s 1.07
add_two / JaXPipe / cuda / PostRev 0.000033856 s 0.000032127000000000006 s 1.05
add_two / JaXPipe / cuda / BothRev 0.000032767999999999995 s 0.000032736 s 1.00
add_two / Jax / cuda / BothRev 0.000034015 s 0.000032191 s 1.06
add_two / HLOOpt / cuda / PreRev 0.000034016 s 0.000032384 s 1.05
add_two / HLOOpt / cuda / PostRev 0.000033728 s 0.000032256 s 1.05
add_two / HLOOpt / cuda / BothRev 0.000033376 s 0.000032384 s 1.03
add_two / PartOpt / cuda / PreRev 0.000033663 s 0.000032255 s 1.04
add_two / PartOpt / cuda / PostRev 0.000033471 s 0.000032384 s 1.03
add_two / PartOpt / cuda / BothRev 0.000032992 s 0.000032032 s 1.03
add_two / IPartOpt / cuda / PreRev 0.000034272 s 0.000033087 s 1.04
add_two / IPartOpt / cuda / PostRev 0.000033631000000000004 s 0.00003152 s 1.07
add_two / IPartOpt / cuda / BothRev 0.000033504 s 0.000032575 s 1.03
add_two / DefOpt / cuda / PreRev 0.000036863 s 0.000032384 s 1.14
add_two / DefOpt / cuda / PostRev 0.000037248000000000007 s 0.000031808000000000004 s 1.17
add_two / DefOpt / cuda / BothRev 0.00003696 s 0.000032448 s 1.14
add_two / IDefOpt / cuda / PreRev 0.000037215 s 0.000032256 s 1.15
add_two / IDefOpt / cuda / PostRev 0.000033856 s 0.00003184 s 1.06
add_two / IDefOpt / cuda / BothRev 0.000033376 s 0.000031776 s 1.05
add_two / JaXPipe / tpu / Primal 0.000001406875 s 0.0000014347750000000002 s 0.98
add_two / Jax / tpu / Primal 0.000001415475 s 0.000001478 s 0.96
add_two / HLOOpt / tpu / Primal 0.000001417075 s 0.0000014378750000000002 s 0.99
add_two / PartOpt / tpu / Primal 0.0000014152000000000002 s 0.0000014793249999999998 s 0.96
add_two / IPartOpt / tpu / Primal 0.000001420475 s 0.0000014442749999999998 s 0.98
add_two / DefOpt / tpu / Primal 0.000001407425 s 0.000001471275 s 0.96
add_two / IDefOpt / tpu / Primal 0.0000014077 s 0.000001440075 s 0.98
add_two / JaXPipe / tpu / Forward 0.000001809375 s 0.00000183065 s 0.99
add_two / Jax / tpu / Forward 0.000001809525 s 0.000001831625 s 0.99
add_two / HLOOpt / tpu / Forward 0.0000018095 s 0.000001825575 s 0.99
add_two / PartOpt / tpu / Forward 0.000001803875 s 0.000001826125 s 0.99
add_two / IPartOpt / tpu / Forward 0.000001818125 s 0.0000018336 s 0.99
add_two / DefOpt / tpu / Forward 0.0000018025 s 0.000001827625 s 0.99
add_two / IDefOpt / tpu / Forward 0.000001816575 s 0.0000018318 s 0.99
add_two / JaXPipe / tpu / PreRev 0.00000281305 s 0.000002840625 s 0.99
add_two / JaXPipe / tpu / PostRev 0.00000273905 s 0.0000027612500000000005 s 0.99
add_two / JaXPipe / tpu / BothRev 0.00000281225 s 0.0000028416 s 0.99
add_two / Jax / tpu / BothRev 0.000002742325 s 0.000002758925 s 0.99
add_two / HLOOpt / tpu / PreRev 0.0000028150750000000003 s 0.000002828875 s 1.00
add_two / HLOOpt / tpu / PostRev 0.000002725775 s 0.0000027550500000000004 s 0.99
add_two / HLOOpt / tpu / BothRev 0.000002800875 s 0.000002834525 s 0.99
add_two / PartOpt / tpu / PreRev 0.000002725625 s 0.0000027605000000000005 s 0.99
add_two / PartOpt / tpu / PostRev 0.00000280215 s 0.000002838175 s 0.99
add_two / PartOpt / tpu / BothRev 0.000002726375 s 0.000002754175 s 0.99
add_two / IPartOpt / tpu / PreRev 0.00000280205 s 0.0000028504000000000003 s 0.98
add_two / IPartOpt / tpu / PostRev 0.000002732825 s 0.000002743325 s 1.00
add_two / IPartOpt / tpu / BothRev 0.000002812025 s 0.000002840525 s 0.99
add_two / DefOpt / tpu / PreRev 0.0000027382250000000004 s 0.0000027612500000000005 s 0.99
add_two / DefOpt / tpu / PostRev 0.00000280335 s 0.000002837125 s 0.99
add_two / DefOpt / tpu / BothRev 0.00000273245 s 0.000002749525 s 0.99
add_two / IDefOpt / tpu / PreRev 0.00000280845 s 0.000002838625 s 0.99
add_two / IDefOpt / tpu / PostRev 0.0000027355 s 0.0000027566 s 0.99
add_two / IDefOpt / tpu / BothRev 0.000002812875 s 0.000002840325 s 0.99
add_two / JaXPipe / cpu / Primal 0.000013203 s 0.000006892639967190917 s 1.92
add_two / Jax / cpu / Primal 0.000021662 s 0.000006933860004210146 s 3.12
add_two / HLOOpt / cpu / Primal 0.000013343 s 0.000007509379993280163 s 1.78
add_two / PartOpt / cpu / Primal 0.000013354 s 0.000006904719966769335 s 1.93
add_two / IPartOpt / cpu / Primal 0.000013261 s 0.000007135020032364992 s 1.86
add_two / DefOpt / cpu / Primal 0.000013046 s 0.000006958459998713806 s 1.87
add_two / IDefOpt / cpu / Primal 0.000013409000000000002 s 0.0000065641800119919935 s 2.04
add_two / JaXPipe / cpu / Forward 0.000018068 s 0.000010514580008020858 s 1.72
add_two / Jax / cpu / Forward 0.000018168 s 0.000010020859963333351 s 1.81
add_two / HLOOpt / cpu / Forward 0.000017971 s 0.000010757099989859853 s 1.67
add_two / PartOpt / cpu / Forward 0.000018651 s 0.000010200499991697142 s 1.83
add_two / IPartOpt / cpu / Forward 0.000018056 s 0.000010306379981557256 s 1.75
add_two / DefOpt / cpu / Forward 0.000018357 s 0.00001038984000842902 s 1.77
add_two / IDefOpt / cpu / Forward 0.000018033 s 0.000010215420024906052 s 1.77
add_two / JaXPipe / cpu / PreRev 0.000023245 s 0.000013841239997418598 s 1.68
add_two / JaXPipe / cpu / PostRev 0.000022908 s 0.000014024039992364124 s 1.63
add_two / JaXPipe / cpu / BothRev 0.000023288 s 0.000014294040029199096 s 1.63
add_two / Jax / cpu / BothRev 0.000023501 s 0.000013598579989775315 s 1.73
add_two / HLOOpt / cpu / PreRev 0.000023166 s 0.00001361135996376106 s 1.70
add_two / HLOOpt / cpu / PostRev 0.000023617 s 0.000016257539991784144 s 1.45
add_two / HLOOpt / cpu / BothRev 0.000023763 s 0.000014164219992380822 s 1.68
add_two / PartOpt / cpu / PreRev 0.000023735 s 0.000013886520009691596 s 1.71
add_two / PartOpt / cpu / PostRev 0.000023074 s 0.00001393359996654908 s 1.66
add_two / PartOpt / cpu / BothRev 0.000023382 s 0.00001496692000728217 s 1.56
add_two / IPartOpt / cpu / PreRev 0.000023674 s 0.000013847800000803544 s 1.71
add_two / IPartOpt / cpu / PostRev 0.000023279 s 0.000014076920015213546 s 1.65
add_two / IPartOpt / cpu / BothRev 0.000023363 s 0.000014492899990727893 s 1.61
add_two / DefOpt / cpu / PreRev 0.000023915 s 0.00001408632003403909 s 1.70
add_two / DefOpt / cpu / PostRev 0.000023623 s 0.000014134579969322658 s 1.67
add_two / DefOpt / cpu / BothRev 0.000023185 s 0.000014213100002962163 s 1.63
add_two / IDefOpt / cpu / PreRev 0.000023761 s 0.000013673720004589996 s 1.74
add_two / IDefOpt / cpu / PostRev 0.000023683 s 0.000014411659985853476 s 1.64
add_two / IDefOpt / cpu / BothRev 0.000023523 s 0.0000148106800406822 s 1.59
cache / JaXPipe / cpu / Primal 0.000006823140001870342 s 0.000006208280001374078 s 1.10
cache / Jax / cpu / Primal 0.000007015360006334959 s 0.0000063180999859469015 s 1.11
cache / HLOOpt / cpu / Primal 0.000007031359973552753 s 0.0000059947600129817145 s 1.17
cache / PartOpt / cpu / Primal 0.000006993879978836049 s 0.000005927159972998197 s 1.18
cache / IPartOpt / cpu / Primal 0.000006802299958508229 s 0.000006284619985308382 s 1.08
cache / DefOpt / cpu / Primal 0.000006829339981777593 s 0.000006352519994834438 s 1.08
cache / IDefOpt / cpu / Primal 0.000006977259990890161 s 0.000006789879980715341 s 1.03
cache / JaXPipe / cpu / Forward 0.000015564619989163475 s 0.000014292039995780216 s 1.09
cache / Jax / cpu / Forward 0.000014708280004924746 s 0.000014286380037447087 s 1.03
cache / HLOOpt / cpu / Forward 0.00001607086001058633 s 0.000016100380007628702 s 1.00
cache / PartOpt / cpu / Forward 0.000015385719980258727 s 0.00001493561997449433 s 1.03
cache / IPartOpt / cpu / Forward 0.000016035159997045413 s 0.000015857279977353754 s 1.01
cache / DefOpt / cpu / Forward 0.000016331719980371417 s 0.00001508203996309021 s 1.08
cache / IDefOpt / cpu / Forward 0.000016240039994954712 s 0.000015270519970727036 s 1.06
cache / JaXPipe / cpu / PreRev 0.000016455719996884 s 0.0000160701800177776 s 1.02
cache / JaXPipe / cpu / PostRev 0.00002167743999962113 s 0.00002165167999010009 s 1.00
cache / JaXPipe / cpu / BothRev 0.000017518139984531444 s 0.000016645660007270636 s 1.05
cache / Jax / cpu / BothRev 0.000020878440018350372 s 0.00002096367997182824 s 1.00
cache / HLOOpt / cpu / PreRev 0.000017311980036538445 s 0.000016826599967316725 s 1.03
cache / HLOOpt / cpu / PostRev 0.00001933374001964694 s 0.000017899080003189738 s 1.08
cache / HLOOpt / cpu / BothRev 0.000016403760037064784 s 0.00001644926000153646 s 1.00
cache / PartOpt / cpu / PreRev 0.000016780160040070767 s 0.00001535305999823322 s 1.09
cache / PartOpt / cpu / PostRev 0.00002121809997333912 s 0.000020587000026353053 s 1.03
cache / PartOpt / cpu / BothRev 0.000017308680016867585 s 0.000015736519981146557 s 1.10
cache / IPartOpt / cpu / PreRev 0.000016950099998211953 s 0.00001527264000287687 s 1.11
cache / IPartOpt / cpu / PostRev 0.00002246466000542569 s 0.00002052838002782664 s 1.09
cache / IPartOpt / cpu / BothRev 0.00001637555998058815 s 0.000015178359981291576 s 1.08
cache / DefOpt / cpu / PreRev 0.000016840039997987332 s 0.000015391139995699633 s 1.09
cache / DefOpt / cpu / PostRev 0.000017693000017970917 s 0.000015205459985736523 s 1.16
cache / DefOpt / cpu / BothRev 0.00001648389999900246 s 0.000015585520031891065 s 1.06
cache / IDefOpt / cpu / PreRev 0.000016591819967288757 s 0.000015139099969019298 s 1.10
cache / IDefOpt / cpu / PostRev 0.00001710182002170768 s 0.000015227000030790805 s 1.12
cache / IDefOpt / cpu / BothRev 0.000016425419962615705 s 0.000015294940003514056 s 1.07
cache / JaXPipe / cuda / Primal 0.000002336 s 0.000002335 s 1.00
cache / Jax / cuda / Primal 0.000002336 s 0.000002336 s 1
cache / HLOOpt / cuda / Primal 0.000002336 s 0.000002273 s 1.03
cache / PartOpt / cuda / Primal 0.000002335 s 0.000002304 s 1.01
cache / IPartOpt / cuda / Primal 0.000002335 s 0.000002335 s 1
cache / DefOpt / cuda / Primal 0.000002335 s 0.000002272 s 1.03
cache / IDefOpt / cuda / Primal 0.000002335 s 0.000002272 s 1.03
cache / JaXPipe / cuda / Forward 0.0000023670000000000004 s 0.000002336 s 1.01
cache / Jax / cuda / Forward 0.0000023670000000000004 s 0.0000023670000000000004 s 1
cache / HLOOpt / cuda / Forward 0.0000023670000000000004 s 0.000002336 s 1.01
cache / PartOpt / cuda / Forward 0.0000023670000000000004 s 0.0000023670000000000004 s 1
cache / IPartOpt / cuda / Forward 0.0000023670000000000004 s 0.000002336 s 1.01
cache / DefOpt / cuda / Forward 0.0000023670000000000004 s 0.000002272 s 1.04
cache / IDefOpt / cuda / Forward 0.0000023670000000000004 s 0.000002336 s 1.01
cache / JaXPipe / cuda / PreRev 0.000010784 s 0.000010687 s 1.01
cache / JaXPipe / cuda / PostRev 0.000010592 s 0.000010847 s 0.98
cache / JaXPipe / cuda / BothRev 0.000010783 s 0.00001072 s 1.01
cache / Jax / cuda / BothRev 0.000010785 s 0.00001104 s 0.98
cache / HLOOpt / cuda / PreRev 0.000013792 s 0.00001312 s 1.05
cache / HLOOpt / cuda / PostRev 0.00001376 s 0.000013088 s 1.05
cache / HLOOpt / cuda / BothRev 0.000013792 s 0.000013088 s 1.05
cache / PartOpt / cuda / PreRev 0.000010784 s 0.000010689 s 1.01
cache / PartOpt / cuda / PostRev 0.000010592 s 0.000010752 s 0.99
cache / PartOpt / cuda / BothRev 0.000010624 s 0.00001088 s 0.98
cache / IPartOpt / cuda / PreRev 0.00001072 s 0.000010752 s 1.00
cache / IPartOpt / cuda / PostRev 0.000010688 s 0.000010464 s 1.02
cache / IPartOpt / cuda / BothRev 0.000010656 s 0.000011968 s 0.89
cache / DefOpt / cuda / PreRev 0.00001072 s 0.000012736 s 0.84
cache / DefOpt / cuda / PostRev 0.000010528 s 0.000010976 s 0.96
cache / DefOpt / cuda / BothRev 0.000010592 s 0.000011007 s 0.96
cache / IDefOpt / cuda / PreRev 0.000010751 s 0.00001072 s 1.00
cache / IDefOpt / cuda / PostRev 0.000010848 s 0.000010432 s 1.04
cache / IDefOpt / cuda / BothRev 0.000010591 s 0.000010272 s 1.03
cache / JaXPipe / tpu / Primal 0.000002475475 s 0.00000247245 s 1.00
cache / Jax / tpu / Primal 0.000002461075 s 0.00000248475 s 0.99
cache / HLOOpt / tpu / Primal 0.00000244885 s 0.0000024782000000000003 s 0.99
cache / PartOpt / tpu / Primal 0.0000024666 s 0.00000247335 s 1.00
cache / IPartOpt / tpu / Primal 0.0000024583 s 0.000002474875 s 0.99
cache / DefOpt / tpu / Primal 0.000002462525 s 0.0000024653 s 1.00
cache / IDefOpt / tpu / Primal 0.000002454775 s 0.00000247745 s 0.99
cache / JaXPipe / tpu / Forward 0.0000035325 s 0.00000353595 s 1.00
cache / Jax / tpu / Forward 0.0000035244 s 0.0000035281 s 1.00
cache / HLOOpt / tpu / Forward 0.0000035349 s 0.00000354365 s 1.00
cache / PartOpt / tpu / Forward 0.0000035322500000000004 s 0.00000353045 s 1.00
cache / IPartOpt / tpu / Forward 0.00000354885 s 0.000003554525 s 1.00
cache / DefOpt / tpu / Forward 0.000003532175 s 0.00000352865 s 1.00
cache / IDefOpt / tpu / Forward 0.0000035453 s 0.000003546525 s 1.00
cache / JaXPipe / tpu / PreRev 0.00000495565 s 0.000004956274999999999 s 1.00
cache / JaXPipe / tpu / PostRev 0.000004975075 s 0.0000049525250000000005 s 1.00
cache / JaXPipe / tpu / BothRev 0.000004974075 s 0.0000049765 s 1.00
cache / Jax / tpu / BothRev 0.000004985425 s 0.00000501655 s 0.99
cache / HLOOpt / tpu / PreRev 0.000003941775 s 0.000003950025 s 1.00
cache / HLOOpt / tpu / PostRev 0.000004128049999999999 s 0.000004126874999999999 s 1.00
cache / HLOOpt / tpu / BothRev 0.000003950775 s 0.0000039442 s 1.00
cache / PartOpt / tpu / PreRev 0.000005000274999999999 s 0.000004969825 s 1.01
cache / PartOpt / tpu / PostRev 0.00000497385 s 0.000004958925 s 1.00
cache / PartOpt / tpu / BothRev 0.000004981725 s 0.0000049805 s 1.00
cache / IPartOpt / tpu / PreRev 0.00000496485 s 0.000004979225000000001 s 1.00
cache / IPartOpt / tpu / PostRev 0.000004971625 s 0.000004973025 s 1.00
cache / IPartOpt / tpu / BothRev 0.000004959749999999999 s 0.0000049724 s 1.00
cache / DefOpt / tpu / PreRev 0.000004961175 s 0.000005000974999999999 s 0.99
cache / DefOpt / tpu / PostRev 0.00000496755 s 0.000004967875 s 1.00
cache / DefOpt / tpu / BothRev 0.000004962675 s 0.000004958125 s 1.00
cache / IDefOpt / tpu / PreRev 0.00000495775 s 0.00000496475 s 1.00
cache / IDefOpt / tpu / PostRev 0.000004974025 s 0.000004959200000000001 s 1.00
cache / IDefOpt / tpu / BothRev 0.0000049764 s 0.0000049626000000000006 s 1.00
cache / JaXPipe / cpu / Primal 0.000012527 s 0.000006208280001374078 s 2.02
cache / Jax / cpu / Primal 0.00001256 s 0.0000063180999859469015 s 1.99
cache / HLOOpt / cpu / Primal 0.000012249 s 0.0000059947600129817145 s 2.04
cache / PartOpt / cpu / Primal 0.000012735 s 0.000005927159972998197 s 2.15
cache / IPartOpt / cpu / Primal 0.000012357 s 0.000006284619985308382 s 1.97
cache / DefOpt / cpu / Primal 0.000012144 s 0.000006352519994834438 s 1.91
cache / IDefOpt / cpu / Primal 0.000012683 s 0.000006789879980715341 s 1.87
cache / JaXPipe / cpu / Forward 0.000017388 s 0.000014292039995780216 s 1.22
cache / Jax / cpu / Forward 0.000017131 s 0.000014286380037447087 s 1.20
cache / HLOOpt / cpu / Forward 0.000017341 s 0.000016100380007628702 s 1.08
cache / PartOpt / cpu / Forward 0.00001734 s 0.00001493561997449433 s 1.16
cache / IPartOpt / cpu / Forward 0.00001712 s 0.000015857279977353754 s 1.08
cache / DefOpt / cpu / Forward 0.000017552 s 0.00001508203996309021 s 1.16
cache / IDefOpt / cpu / Forward 0.000017487 s 0.000015270519970727036 s 1.15
cache / JaXPipe / cpu / PreRev 0.000017973 s 0.0000160701800177776 s 1.12
cache / JaXPipe / cpu / PostRev 0.000021131 s 0.00002165167999010009 s 0.98
cache / JaXPipe / cpu / BothRev 0.000018906 s 0.000016645660007270636 s 1.14
cache / Jax / cpu / BothRev 0.000021291 s 0.00002096367997182824 s 1.02
cache / HLOOpt / cpu / PreRev 0.00001825 s 0.000016826599967316725 s 1.08
cache / HLOOpt / cpu / PostRev 0.00001869 s 0.000017899080003189738 s 1.04
cache / HLOOpt / cpu / BothRev 0.000017788000000000003 s 0.00001644926000153646 s 1.08
cache / PartOpt / cpu / PreRev 0.000018127 s 0.00001535305999823322 s 1.18
cache / PartOpt / cpu / PostRev 0.00002188 s 0.000020587000026353053 s 1.06
cache / PartOpt / cpu / BothRev 0.000018076 s 0.000015736519981146557 s 1.15
cache / IPartOpt / cpu / PreRev 0.000018655 s 0.00001527264000287687 s 1.22
cache / IPartOpt / cpu / PostRev 0.000021406 s 0.00002052838002782664 s 1.04
cache / IPartOpt / cpu / BothRev 0.000018163 s 0.000015178359981291576 s 1.20
cache / DefOpt / cpu / PreRev 0.000031324 s 0.000015391139995699633 s 2.04
cache / DefOpt / cpu / PostRev 0.000017318 s 0.000015205459985736523 s 1.14
cache / DefOpt / cpu / BothRev 0.000018135 s 0.000015585520031891065 s 1.16
cache / IDefOpt / cpu / PreRev 0.000018685 s 0.000015139099969019298 s 1.23
cache / IDefOpt / cpu / PostRev 0.000018271 s 0.000015227000030790805 s 1.20
cache / IDefOpt / cpu / BothRev 0.000018023 s 0.000015294940003514056 s 1.18
Concat / JaXPipe / cpu / Primal 0.00000766007993661333 s 0.00000680945997373783 s 1.12
Concat / Jax / cpu / Primal 0.000007366859981630114 s 0.00000685994000377832 s 1.07
Concat / HLOOpt / cpu / Primal 0.000007768140003463486 s 0.000007261000018843333 s 1.07
Concat / PartOpt / cpu / Primal 0.000007408859974020743 s 0.000006560999991052086 s 1.13
Concat / IPartOpt / cpu / Primal 0.000007114780028132373 s 0.00000664078000227164 s 1.07
Concat / DefOpt / cpu / Primal 0.000007462180037691723 s 0.000006791279938624939 s 1.10
Concat / IDefOpt / cpu / Primal 0.000007165919996623416 s 0.000006834340010755113 s 1.05
Concat / JaXPipe / cpu / Forward 0.00001086562000637059 s 0.000010419480004202342 s 1.04
Concat / Jax / cpu / Forward 0.000011015580012099236 s 0.000009747220028657466 s 1.13
Concat / HLOOpt / cpu / Forward 0.00001093417998163204 s 0.000010163860006286995 s 1.08
Concat / PartOpt / cpu / Forward 0.000011337880014252733 s 0.000010052399966298254 s 1.13
Concat / IPartOpt / cpu / Forward 0.000011647180008367286 s 0.000010310499992556289 s 1.13
Concat / DefOpt / cpu / Forward 0.000011288460009382106 s 0.000010081060054289992 s 1.12
Concat / IDefOpt / cpu / Forward 0.000011130779994346084 s 0.000009757579982760945 s 1.14
Concat / JaXPipe / cpu / PreRev 0.000013252600010673632 s 0.000011575919988899842 s 1.14
Concat / JaXPipe / cpu / PostRev 0.000013369819989748064 s 0.000011765639956138327 s 1.14
Concat / JaXPipe / cpu / BothRev 0.000013727739997193568 s 0.000011440880025475052 s 1.20
Concat / Jax / cpu / BothRev 0.000012901099980808794 s 0.000011740379995899275 s 1.10
Concat / HLOOpt / cpu / PreRev 0.000014216339977792811 s 0.000012096380041839438 s 1.18
Concat / HLOOpt / cpu / PostRev 0.000014944299955459428 s 0.000013298899993969826 s 1.12
Concat / HLOOpt / cpu / BothRev 0.000013100780015520288 s 0.000011380740015738411 s 1.15
Concat / PartOpt / cpu / PreRev 0.000013416759984465896 s 0.000011581600019781036 s 1.16
Concat / PartOpt / cpu / PostRev 0.00001248861998647044 s 0.000011830000012196251 s 1.06
Concat / PartOpt / cpu / BothRev 0.000013253720017019077 s 0.00001163582000117458 s 1.14
Concat / IPartOpt / cpu / PreRev 0.000012622819995158351 s 0.000011728999961633237 s 1.08
Concat / IPartOpt / cpu / PostRev 0.000012947659988640226 s 0.000011372859962648362 s 1.14
Concat / IPartOpt / cpu / BothRev 0.000013015000004088506 s 0.00001142194001658936 s 1.14
Concat / DefOpt / cpu / PreRev 0.0000126365599953715 s 0.000011049559971070266 s 1.14
Concat / DefOpt / cpu / PostRev 0.00001276639998650353 s 0.000011627620033323182 s 1.10
Concat / DefOpt / cpu / BothRev 0.000013241280021247805 s 0.00001125859998865053 s 1.18
Concat / IDefOpt / cpu / PreRev 0.000013696740015802789 s 0.00001125570005569898 s 1.22
Concat / IDefOpt / cpu / PostRev 0.000012427799965735176 s 0.000011331400010021751 s 1.10
Concat / IDefOpt / cpu / BothRev 0.000012435040007403586 s 0.000011248080018049222 s 1.11
Concat / JaXPipe / cuda / Primal 0.000002464 s 0.0000019200000000000003 s 1.28
Concat / Jax / cuda / Primal 0.000002464 s 0.0000019200000000000003 s 1.28
Concat / HLOOpt / cuda / Primal 0.000002464 s 0.000001919 s 1.28
Concat / PartOpt / cuda / Primal 0.000002464 s 0.000001888 s 1.31
Concat / IPartOpt / cuda / Primal 0.000002463 s 0.0000019200000000000003 s 1.28
Concat / DefOpt / cuda / Primal 0.000002463 s 0.0000019200000000000003 s 1.28
Concat / IDefOpt / cuda / Primal 0.000002463 s 0.0000019200000000000003 s 1.28
Concat / JaXPipe / cuda / Forward 0.000010912 s 0.00001008 s 1.08
Concat / Jax / cuda / Forward 0.00001104 s 0.000009984 s 1.11
Concat / HLOOpt / cuda / Forward 0.000010625 s 0.000010048 s 1.06
Concat / PartOpt / cuda / Forward 0.000010688 s 0.000010145 s 1.05
Concat / IPartOpt / cuda / Forward 0.00001072 s 0.000009984 s 1.07
Concat / DefOpt / cuda / Forward 0.000010688 s 0.000009984 s 1.07
Concat / IDefOpt / cuda / Forward 0.000010016 s 0.000009985 s 1.00
Concat / JaXPipe / cuda / PreRev 0.000017247999999999998 s 0.000016607 s 1.04
Concat / JaXPipe / cuda / PostRev 0.000017152 s 0.000015552 s 1.10
Concat / JaXPipe / cuda / BothRev 0.000016864 s 0.00001648 s 1.02
Concat / Jax / cuda / BothRev 0.000017344 s 0.000016448000000000002 s 1.05
Concat / HLOOpt / cuda / PreRev 0.000017151 s 0.000016255999999999998 s 1.06
Concat / HLOOpt / cuda / PostRev 0.00001664 s 0.000016063999999999997 s 1.04
Concat / HLOOpt / cuda / BothRev 0.000016607 s 0.000016224 s 1.02
Concat / PartOpt / cuda / PreRev 0.000017056 s 0.000016832 s 1.01
Concat / PartOpt / cuda / PostRev 0.000016927000000000002 s 0.000016192 s 1.05
Concat / PartOpt / cuda / BothRev 0.000017343 s 0.00001648 s 1.05
Concat / IPartOpt / cuda / PreRev 0.000017216 s 0.000016576000000000002 s 1.04
Concat / IPartOpt / cuda / PostRev 0.000016896000000000002 s 0.00001664 s 1.02
Concat / IPartOpt / cuda / BothRev 0.000016992 s 0.000016383999999999998 s 1.04
Concat / DefOpt / cuda / PreRev 0.0000168 s 0.00001696 s 0.99
Concat / DefOpt / cuda / PostRev 0.000016863 s 0.000016319 s 1.03
Concat / DefOpt / cuda / BothRev 0.000016447 s 0.00001664 s 0.99
Concat / IDefOpt / cuda / PreRev 0.000016958999999999998 s 0.00001648 s 1.03
Concat / IDefOpt / cuda / PostRev 0.000016896000000000002 s 0.00001552 s 1.09
Concat / IDefOpt / cuda / BothRev 0.000016927999999999998 s 0.00001616 s 1.05
Concat / JaXPipe / tpu / Primal 0.0000014848000000000002 s 0.000001527525 s 0.97
Concat / Jax / tpu / Primal 0.0000014771 s 0.000001528825 s 0.97
Concat / HLOOpt / tpu / Primal 0.0000014774 s 0.00000152625 s 0.97
Concat / PartOpt / tpu / Primal 0.0000014781 s 0.0000015351 s 0.96
Concat / IPartOpt / tpu / Primal 0.000001488025 s 0.000001526125 s 0.98
Concat / DefOpt / tpu / Primal 0.000001468525 s 0.0000015325 s 0.96
Concat / IDefOpt / tpu / Primal 0.0000014741000000000002 s 0.000001542 s 0.96
Concat / JaXPipe / tpu / Forward 0.0000015559249999999998 s 0.0000015702250000000005 s 0.99
Concat / Jax / tpu / Forward 0.0000015204500000000002 s 0.000001549 s 0.98
Concat / HLOOpt / tpu / Forward 0.0000015479999999999998 s 0.00000158455 s 0.98
Concat / PartOpt / tpu / Forward 0.00000151865 s 0.00000155495 s 0.98
Concat / IPartOpt / tpu / Forward 0.000001539225 s 0.000001570625 s 0.98
Concat / DefOpt / tpu / Forward 0.00000152255 s 0.0000015566 s 0.98
Concat / IDefOpt / tpu / Forward 0.000001542075 s 0.0000015707750000000002 s 0.98
Concat / JaXPipe / tpu / PreRev 0.0000019563 s 0.000001993525 s 0.98
Concat / JaXPipe / tpu / PostRev 0.0000020299250000000003 s 0.00000208665 s 0.97
Concat / JaXPipe / tpu / BothRev 0.0000019493750000000003 s 0.0000019998250000000004 s 0.97
Concat / Jax / tpu / BothRev 0.000002013575 s 0.000002072025 s 0.97
Concat / HLOOpt / tpu / PreRev 0.00000194715 s 0.0000019958000000000004 s 0.98
Concat / HLOOpt / tpu / PostRev 0.000002015875 s 0.00000206755 s 0.98
Concat / HLOOpt / tpu / BothRev 0.0000019507250000000003 s 0.0000020007 s 0.98
Concat / PartOpt / tpu / PreRev 0.00000201765 s 0.0000020732 s 0.97
Concat / PartOpt / tpu / PostRev 0.0000019563250000000003 s 0.0000019972 s 0.98
Concat / PartOpt / tpu / BothRev 0.0000020297 s 0.000002075475 s 0.98
Concat / IPartOpt / tpu / PreRev 0.00000195545 s 0.00000199595 s 0.98
Concat / IPartOpt / tpu / PostRev 0.00000201155 s 0.00000207615 s 0.97
Concat / IPartOpt / tpu / BothRev 0.00000195815 s 0.0000020014 s 0.98
Concat / DefOpt / tpu / PreRev 0.0000020121 s 0.000002076375 s 0.97
Concat / DefOpt / tpu / PostRev 0.0000019502 s 0.000001996825 s 0.98
Concat / DefOpt / tpu / BothRev 0.000002019125 s 0.0000020678 s 0.98
Concat / IDefOpt / tpu / PreRev 0.0000019554 s 0.000002002125 s 0.98
Concat / IDefOpt / tpu / PostRev 0.000002016625 s 0.000002064275 s 0.98
Concat / IDefOpt / tpu / BothRev 0.0000019455 s 0.000001992175 s 0.98
Concat / JaXPipe / cpu / Primal 0.00001297 s 0.00000680945997373783 s 1.90
Concat / Jax / cpu / Primal 0.000012602 s 0.00000685994000377832 s 1.84
Concat / HLOOpt / cpu / Primal 0.000012263 s 0.000007261000018843333 s 1.69
Concat / PartOpt / cpu / Primal 0.000012615 s 0.000006560999991052086 s 1.92
Concat / IPartOpt / cpu / Primal 0.000012648 s 0.00000664078000227164 s 1.90
Concat / DefOpt / cpu / Primal 0.000012922 s 0.000006791279938624939 s 1.90
Concat / IDefOpt / cpu / Primal 0.000012429 s 0.000006834340010755113 s 1.82
Concat / JaXPipe / cpu / Forward 0.000017598 s 0.000010419480004202342 s 1.69
Concat / Jax / cpu / Forward 0.000017676 s 0.000009747220028657466 s 1.81
Concat / HLOOpt / cpu / Forward 0.000018118 s 0.000010163860006286995 s 1.78
Concat / PartOpt / cpu / Forward 0.000017967 s 0.000010052399966298254 s 1.79
Concat / IPartOpt / cpu / Forward 0.000017768 s 0.000010310499992556289 s 1.72
Concat / DefOpt / cpu / Forward 0.000017868 s 0.000010081060054289992 s 1.77
Concat / IDefOpt / cpu / Forward 0.000017763999999999998 s 0.000009757579982760945 s 1.82
Concat / JaXPipe / cpu / PreRev 0.000020428 s 0.000011575919988899842 s 1.76
Concat / JaXPipe / cpu / PostRev 0.000019884 s 0.000011765639956138327 s 1.69
Concat / JaXPipe / cpu / BothRev 0.000020025 s 0.000011440880025475052 s 1.75
Concat / Jax / cpu / BothRev 0.000020088 s 0.000011740379995899275 s 1.71
Concat / HLOOpt / cpu / PreRev 0.00001991 s 0.000012096380041839438 s 1.65
Concat / HLOOpt / cpu / PostRev 0.000020189 s 0.000013298899993969826 s 1.52
Concat / HLOOpt / cpu / BothRev 0.000020349 s 0.000011380740015738411 s 1.79
Concat / PartOpt / cpu / PreRev 0.00002049 s 0.000011581600019781036 s 1.77
Concat / PartOpt / cpu / PostRev 0.000020334 s 0.000011830000012196251 s 1.72
Concat / PartOpt / cpu / BothRev 0.000019354 s 0.00001163582000117458 s 1.66
Concat / IPartOpt / cpu / PreRev 0.000020246 s 0.000011728999961633237 s 1.73
Concat / IPartOpt / cpu / PostRev 0.00002005 s 0.000011372859962648362 s 1.76
Concat / IPartOpt / cpu / BothRev 0.000020425 s 0.00001142194001658936 s 1.79
Concat / DefOpt / cpu / PreRev 0.000020366 s 0.000011049559971070266 s 1.84
Concat / DefOpt / cpu / PostRev 0.000019418000000000003 s 0.000011627620033323182 s 1.67
Concat / DefOpt / cpu / BothRev 0.00002048 s 0.00001125859998865053 s 1.82
Concat / IDefOpt / cpu / PreRev 0.000019896 s 0.00001125570005569898 s 1.77
Concat / IDefOpt / cpu / PostRev 0.000020015 s 0.000011331400010021751 s 1.77
Concat / IDefOpt / cpu / BothRev 0.000020316 s 0.000011248080018049222 s 1.81
const_scatter / JaXPipe / cpu / Primal 0.000007343920005951077 s 0.000006517819956570747 s 1.13
const_scatter / Jax / cpu / Primal 0.000007374120013992069 s 0.000006735540018780739 s 1.09
const_scatter / HLOOpt / cpu / Primal 0.000008084000010057934 s 0.000007361260004472569 s 1.10
const_scatter / PartOpt / cpu / Primal 0.000007260040047185612 s 0.000006297840009210632 s 1.15
const_scatter / IPartOpt / cpu / Primal 0.000007671040002605878 s 0.0000062833599986333865 s 1.22
const_scatter / DefOpt / cpu / Primal 0.000007151119989430299 s 0.0000069139400511630815 s 1.03
const_scatter / IDefOpt / cpu / Primal 0.000007917839993751841 s 0.0000069470600101340094 s 1.14
const_scatter / JaXPipe / cpu / Forward 0.000012128340003982884 s 0.000010408939997432754 s 1.17
const_scatter / Jax / cpu / Forward 0.00001087952001398662 s 0.000009712399996715247 s 1.12
const_scatter / HLOOpt / cpu / Forward 0.000012555019984574756 s 0.00001093629999559198 s 1.15
const_scatter / PartOpt / cpu / Forward 0.000011797820015999604 s 0.00001031351998790342 s 1.14
const_scatter / IPartOpt / cpu / Forward 0.000012081619997843518 s 0.0000109735800378985 s 1.10
const_scatter / DefOpt / cpu / Forward 0.00001172765998489922 s 0.000010581439992165544 s 1.11
const_scatter / IDefOpt / cpu / Forward 0.000011968500002694782 s 0.000010888860024351744 s 1.10
const_scatter / JaXPipe / cpu / PreRev 0.0002907639799923 s 0.0002880199600076 s 1.01
const_scatter / JaXPipe / cpu / PostRev 0.0002840496599947 s 0.0002795422999861 s 1.02
const_scatter / JaXPipe / cpu / BothRev 0.0002871555400543 s 0.0002823235399318 s 1.02
const_scatter / Jax / cpu / BothRev 0.0002829453799495 s 0.0002785030800077 s 1.02
const_scatter / HLOOpt / cpu / PreRev 0.0002926128000399 s 0.0002912244199433 s 1.00
const_scatter / HLOOpt / cpu / PostRev 0.0002872863999891 s 0.0002848686400193 s 1.01
const_scatter / HLOOpt / cpu / BothRev 0.0002878475799843 s 0.0002826136400199 s 1.02
const_scatter / PartOpt / cpu / PreRev 0.0002966800000376 s 0.0002812760599772 s 1.05
const_scatter / PartOpt / cpu / PostRev 0.000282708299992 s 0.0002809123799943 s 1.01
const_scatter / PartOpt / cpu / BothRev 0.0002864107800269 s 0.0002803485599906 s 1.02
const_scatter / IPartOpt / cpu / PreRev 0.0002875837799729 s 0.0002837234400067 s 1.01
const_scatter / IPartOpt / cpu / PostRev 0.0002850183599912 s 0.0002823269600412 s 1.01
const_scatter / IPartOpt / cpu / BothRev 0.00028588965999 s 0.0002820849200179 s 1.01
const_scatter / DefOpt / cpu / PreRev 0.0002866056799757 s 0.0002825540000412 s 1.01
const_scatter / DefOpt / cpu / PostRev 0.0002863893000176 s 0.000283008159995 s 1.01
const_scatter / DefOpt / cpu / BothRev 0.0002865980399747 s 0.0002790038200146 s 1.03
const_scatter / IDefOpt / cpu / PreRev 0.0002884569399975 s 0.0002833644399834 s 1.02
const_scatter / IDefOpt / cpu / PostRev 0.0002854432199728 s 0.000283859480005 s 1.01
const_scatter / IDefOpt / cpu / BothRev 0.0002857488000063 s 0.0002820218600299 s 1.01
const_scatter / JaXPipe / cuda / Primal 0.000002463 s 0.000001887 s 1.31
const_scatter / Jax / cuda / Primal 0.000002463 s 0.000001887 s 1.31
const_scatter / HLOOpt / cuda / Primal 0.000002463 s 0.000001887 s 1.31
const_scatter / PartOpt / cuda / Primal 0.000002463 s 0.000001887 s 1.31
const_scatter / IPartOpt / cuda / Primal 0.000002463 s 0.000001887 s 1.31
const_scatter / DefOpt / cuda / Primal 0.000002463 s 0.000001887 s 1.31
const_scatter / IDefOpt / cuda / Primal 0.000002463 s 0.000001887 s 1.31
const_scatter / JaXPipe / cuda / Forward 0.000010592 s 0.000010176 s 1.04
const_scatter / Jax / cuda / Forward 0.000010496 s 0.000009759 s 1.08
const_scatter / HLOOpt / cuda / Forward 0.000010336 s 0.000010048 s 1.03
const_scatter / PartOpt / cuda / Forward 0.000010368 s 0.000010048 s 1.03
const_scatter / IPartOpt / cuda / Forward 0.000010496 s 0.000009856 s 1.06
const_scatter / DefOpt / cuda / Forward 0.000010784 s 0.000009953 s 1.08
const_scatter / IDefOpt / cuda / Forward 0.00001056 s 0.000010016 s 1.05
const_scatter / JaXPipe / cuda / PreRev 0.000017024 s 0.000016255999999999998 s 1.05
const_scatter / JaXPipe / cuda / PostRev 0.0000168 s 0.000016255999999999998 s 1.03
const_scatter / JaXPipe / cuda / BothRev 0.000016768000000000003 s 0.000016063000000000002 s 1.04
const_scatter / Jax / cuda / BothRev 0.000016768000000000003 s 0.000016192 s 1.04
const_scatter / HLOOpt / cuda / PreRev 0.000017344 s 0.000016608 s 1.04
const_scatter / HLOOpt / cuda / PostRev 0.000017312 s 0.000016255999999999998 s 1.06
const_scatter / HLOOpt / cuda / BothRev 0.0000168 s 0.000015648 s 1.07
const_scatter / PartOpt / cuda / PreRev 0.000017312 s 0.000016255999999999998 s 1.06
const_scatter / PartOpt / cuda / PostRev 0.000016705 s 0.000016927999999999998 s 0.99
const_scatter / PartOpt / cuda / BothRev 0.00001728 s 0.000015712 s 1.10
const_scatter / IPartOpt / cuda / PreRev 0.000016448000000000002 s 0.00001664 s 0.99
const_scatter / IPartOpt / cuda / PostRev 0.0000168 s 0.000016255999999999998 s 1.03
const_scatter / IPartOpt / cuda / BothRev 0.000016416 s 0.000016192 s 1.01
const_scatter / DefOpt / cuda / PreRev 0.000017855 s 0.00001632 s 1.09
const_scatter / DefOpt / cuda / PostRev 0.000016832 s 0.000016479 s 1.02
const_scatter / DefOpt / cuda / BothRev 0.000016672 s 0.000016128 s 1.03
const_scatter / IDefOpt / cuda / PreRev 0.00001696 s 0.000016832 s 1.01
const_scatter / IDefOpt / cuda / PostRev 0.000016608 s 0.000016192 s 1.03
const_scatter / IDefOpt / cuda / BothRev 0.000016704 s 0.000016416 s 1.02
const_scatter / JaXPipe / tpu / Primal 0.0000037958 s 0.000003786025 s 1.00
const_scatter / Jax / tpu / Primal 0.0000038175750000000005 s 0.000003809375 s 1.00
const_scatter / HLOOpt / tpu / Primal 0.000003802025 s 0.0000037999 s 1.00
const_scatter / PartOpt / tpu / Primal 0.000003802875 s 0.000003804 s 1.00
const_scatter / IPartOpt / tpu / Primal 0.000003777275 s 0.000003802075 s 0.99
const_scatter / DefOpt / tpu / Primal 0.000003794600000000001 s 0.000003828525 s 0.99
const_scatter / IDefOpt / tpu / Primal 0.0000037942 s 0.00000379765 s 1.00
const_scatter / JaXPipe / tpu / Forward 0.00000647495 s 0.000006461425 s 1.00
const_scatter / Jax / tpu / Forward 0.0000064932 s 0.000006504925 s 1.00
const_scatter / HLOOpt / tpu / Forward 0.000006464925 s 0.000006455725 s 1.00
const_scatter / PartOpt / tpu / Forward 0.000006474825 s 0.000006500375 s 1.00
const_scatter / IPartOpt / tpu / Forward 0.0000064609 s 0.000006466675 s 1.00
const_scatter / DefOpt / tpu / Forward 0.000006467300000000001 s 0.00000647215 s 1.00
const_scatter / IDefOpt / tpu / Forward 0.0000064622000000000006 s 0.000006480924999999999 s 1.00
const_scatter / JaXPipe / tpu / PreRev 0.000006600775 s 0.000006688075 s 0.99
const_scatter / JaXPipe / tpu / PostRev 0.000006606049999999999 s 0.000006672975 s 0.99
const_scatter / JaXPipe / tpu / BothRev 0.000006600325000000001 s 0.000006678150000000001 s 0.99
const_scatter / Jax / tpu / BothRev 0.0000066186 s 0.000006678525 s 0.99
const_scatter / HLOOpt / tpu / PreRev 0.0000066067 s 0.0000066666 s 0.99
const_scatter / HLOOpt / tpu / PostRev 0.000006617525 s 0.000006659000000000001 s 0.99
const_scatter / HLOOpt / tpu / BothRev 0.000006608 s 0.000006699475 s 0.99
const_scatter / PartOpt / tpu / PreRev 0.000006596475000000001 s 0.000006648425000000001 s 0.99
const_scatter / PartOpt / tpu / PostRev 0.000006619125 s 0.000006682625 s 0.99
const_scatter / PartOpt / tpu / BothRev 0.000006614825 s 0.00000666005 s 0.99
const_scatter / IPartOpt / tpu / PreRev 0.000006594575 s 0.000006676075 s 0.99
const_scatter / IPartOpt / tpu / PostRev 0.0000066035500000000005 s 0.00000664445 s 0.99
const_scatter / IPartOpt / tpu / BothRev 0.0000065984 s 0.000006668 s 0.99
const_scatter / DefOpt / tpu / PreRev 0.0000066193 s 0.0000066681 s 0.99
const_scatter / DefOpt / tpu / PostRev 0.00000660375 s 0.000006650825 s 0.99
const_scatter / DefOpt / tpu / BothRev 0.0000066391000000000005 s 0.000006653649999999999 s 1.00
const_scatter / IDefOpt / tpu / PreRev 0.0000065939250000000006 s 0.0000066552 s 0.99
const_scatter / IDefOpt / tpu / PostRev 0.000006632625 s 0.000006674099999999999 s 0.99
const_scatter / IDefOpt / tpu / BothRev 0.000006580975 s 0.0000066716 s 0.99
const_scatter / JaXPipe / cpu / Primal 0.000012611 s 0.000006517819956570747 s 1.93
const_scatter / Jax / cpu / Primal 0.000012529 s 0.000006735540018780739 s 1.86
const_scatter / HLOOpt / cpu / Primal 0.000013558 s 0.000007361260004472569 s 1.84
const_scatter / PartOpt / cpu / Primal 0.000012284 s 0.000006297840009210632 s 1.95
const_scatter / IPartOpt / cpu / Primal 0.00001231 s 0.0000062833599986333865 s 1.96
const_scatter / DefOpt / cpu / Primal 0.000013235 s 0.0000069139400511630815 s 1.91
const_scatter / IDefOpt / cpu / Primal 0.000013303 s 0.0000069470600101340094 s 1.91
const_scatter / JaXPipe / cpu / Forward 0.000017927 s 0.000010408939997432754 s 1.72
const_scatter / Jax / cpu / Forward 0.000016427 s 0.000009712399996715247 s 1.69
const_scatter / HLOOpt / cpu / Forward 0.000018229 s 0.00001093629999559198 s 1.67
const_scatter / PartOpt / cpu / Forward 0.000018357 s 0.00001031351998790342 s 1.78
const_scatter / IPartOpt / cpu / Forward 0.000018566 s 0.0000109735800378985 s 1.69
const_scatter / DefOpt / cpu / Forward 0.00001812 s 0.000010581439992165544 s 1.71
const_scatter / IDefOpt / cpu / Forward 0.000018053 s 0.000010888860024351744 s 1.66
const_scatter / JaXPipe / cpu / PreRev 0.000526605 s 0.0002880199600076 s 1.83
const_scatter / JaXPipe / cpu / PostRev 0.000497574 s 0.0002795422999861 s 1.78
const_scatter / JaXPipe / cpu / BothRev 0.000502461 s 0.0002823235399318 s 1.78
const_scatter / Jax / cpu / BothRev 0.000531197 s 0.0002785030800077 s 1.91
const_scatter / HLOOpt / cpu / PreRev 0.0004939 s 0.0002912244199433 s 1.70
const_scatter / HLOOpt / cpu / PostRev 0.000491768 s 0.0002848686400193 s 1.73
const_scatter / HLOOpt / cpu / BothRev 0.000522844 s 0.0002826136400199 s 1.85
const_scatter / PartOpt / cpu / PreRev 0.000525233 s 0.0002812760599772 s 1.87
const_scatter / PartOpt / cpu / PostRev 0.00052519 s 0.0002809123799943 s 1.87
const_scatter / PartOpt / cpu / BothRev 0.000517604 s 0.0002803485599906 s 1.85
const_scatter / IPartOpt / cpu / PreRev 0.000498367 s 0.0002837234400067 s 1.76
const_scatter / IPartOpt / cpu / PostRev 0.000501546 s 0.0002823269600412 s 1.78
const_scatter / IPartOpt / cpu / BothRev 0.000496307 s 0.0002820849200179 s 1.76
const_scatter / DefOpt / cpu / PreRev 0.000508057 s 0.0002825540000412 s 1.80
const_scatter / DefOpt / cpu / PostRev 0.000512242 s 0.000283008159995 s 1.81
const_scatter / DefOpt / cpu / BothRev 0.000492891 s 0.0002790038200146 s 1.77
const_scatter / IDefOpt / cpu / PreRev 0.00050017 s 0.0002833644399834 s 1.77
const_scatter / IDefOpt / cpu / PostRev 0.0004991419999999 s 0.000283859480005 s 1.76
const_scatter / IDefOpt / cpu / BothRev 0.00049004 s 0.0002820218600299 s 1.74
GenDot / JaXPipe / cpu / Primal 0.000008898860005501775 s 0.000006996919955781777 s 1.27
GenDot / Jax / cpu / Primal 0.000008550899965484859 s 0.000007466960050805938 s 1.15
GenDot / HLOOpt / cpu / Primal 0.000008346540007551083 s 0.000007573020011477638 s 1.10
GenDot / PartOpt / cpu / Primal 0.000008447200016235002 s 0.00000727000000551925 s 1.16
GenDot / IPartOpt / cpu / Primal 0.000008423919989581919 s 0.000007262939989232109 s 1.16
GenDot / DefOpt / cpu / Primal 0.000008014699997147545 s 0.000007189280004240572 s 1.11
GenDot / IDefOpt / cpu / Primal 0.000008033539979805936 s 0.000006937339976502699 s 1.16
GenDot / JaXPipe / cpu / Forward 0.000012085880007361991 s 0.00001109208003981621 s 1.09
GenDot / Jax / cpu / Forward 0.000011434980024205288 s 0.000010694279944800656 s 1.07
GenDot / HLOOpt / cpu / Forward 0.000012270239958525051 s 0.000011425719985709292 s 1.07
GenDot / PartOpt / cpu / Forward 0.0000122045199987042 s 0.000011010379994331744 s 1.11
GenDot / IPartOpt / cpu / Forward 0.000012632680027309108 s 0.000011673500002871151 s 1.08
GenDot / DefOpt / cpu / Forward 0.000011274360012976103 s 0.000011217820037927595 s 1.01
GenDot / IDefOpt / cpu / Forward 0.00001214182002513553 s 0.000011160699987158296 s 1.09
GenDot / JaXPipe / cpu / PreRev 0.00001204234000397264 s 0.000011052339996240334 s 1.09
GenDot / JaXPipe / cpu / PostRev 0.00001146877999417484 s 0.000010355519953009206 s 1.11
GenDot / JaXPipe / cpu / BothRev 0.000012770240027748514 s 0.000011711480019584996 s 1.09
GenDot / Jax / cpu / BothRev 0.000011995019976893672 s 0.000010143320014321944 s 1.18
GenDot / HLOOpt / cpu / PreRev 0.000012809980016754708 s 0.00001173204003862338 s 1.09
GenDot / HLOOpt / cpu / PostRev 0.00001452287998290558 s 0.000013266199948702706 s 1.09
GenDot / HLOOpt / cpu / BothRev 0.00001186527999379905 s 0.00001131287996940955 s 1.05
GenDot / PartOpt / cpu / PreRev 0.000012013880004815291 s 0.00001143761996900139 s 1.05
GenDot / PartOpt / cpu / PostRev 0.00001179478002086398 s 0.000010051760045826088 s 1.17
GenDot / PartOpt / cpu / BothRev 0.00001281204000406433 s 0.00001174368006104487 s 1.09
GenDot / IPartOpt / cpu / PreRev 0.00001247651998710353 s 0.000011260359970037826 s 1.11
GenDot / IPartOpt / cpu / PostRev 0.00001132471998971596 s 0.000009902019992296118 s 1.14
GenDot / IPartOpt / cpu / BothRev 0.000012123359983888804 s 0.000011261819963692687 s 1.08
GenDot / DefOpt / cpu / PreRev 0.0000117534599667124 s 0.000011598100036280811 s 1.01
GenDot / DefOpt / cpu / PostRev 0.000012987260006411817 s 0.00001124961999266816 s 1.15
GenDot / DefOpt / cpu / BothRev 0.000011985540004388896 s 0.000011092819977420732 s 1.08
GenDot / IDefOpt / cpu / PreRev 0.00001243700002305559 s 0.000011347700010446716 s 1.10
GenDot / IDefOpt / cpu / PostRev 0.000013013000034334254 s 0.000012072240006091306 s 1.08
GenDot / IDefOpt / cpu / BothRev 0.00001225334002810996 s 0.00001116171999456128 s 1.10
GenDot / JaXPipe / cuda / Primal 0.00000256 s 0.000002015 s 1.27
GenDot / Jax / cuda / Primal 0.000002559 s 0.000002015 s 1.27
GenDot / HLOOpt / cuda / Primal 0.000002527 s 0.000001984 s 1.27
GenDot / PartOpt / cuda / Primal 0.00000256 s 0.000002015 s 1.27
GenDot / IPartOpt / cuda / Primal 0.00000256 s 0.000002016 s 1.27
GenDot / DefOpt / cuda / Primal 0.000002559 s 0.000001984 s 1.29
GenDot / IDefOpt / cuda / Primal 0.000002559 s 0.000001984 s 1.29
GenDot / JaXPipe / cuda / Forward 0.00001072 s 0.000010111 s 1.06
GenDot / Jax / cuda / Forward 0.000010784 s 0.000010048 s 1.07
GenDot / HLOOpt / cuda / Forward 0.000010752 s 0.000010912 s 0.99
GenDot / PartOpt / cuda / Forward 0.000011008 s 0.000010176 s 1.08
GenDot / IPartOpt / cuda / Forward 0.000010656 s 0.000009984 s 1.07
GenDot / DefOpt / cuda / Forward 0.000010784 s 0.000010144 s 1.06
GenDot / IDefOpt / cuda / Forward 0.000010816 s 0.000009792 s 1.10
GenDot / JaXPipe / cuda / PreRev 0.000010464 s 0.000010144 s 1.03
GenDot / JaXPipe / cuda / PostRev 0.00001088 s 0.000010624 s 1.02
GenDot / JaXPipe / cuda / BothRev 0.000010623 s 0.000009824 s 1.08
GenDot / Jax / cuda / BothRev 0.000010944 s 0.00001008 s 1.09
GenDot / HLOOpt / cuda / PreRev 0.000010304 s 0.000010144 s 1.02
GenDot / HLOOpt / cuda / PostRev 0.000010688 s 0.000010048 s 1.06
GenDot / HLOOpt / cuda / BothRev 0.000010816 s 0.000010016 s 1.08
GenDot / PartOpt / cuda / PreRev 0.000010623 s 0.000010368 s 1.02
GenDot / PartOpt / cuda / PostRev 0.000010495 s 0.00000976 s 1.08
GenDot / PartOpt / cuda / BothRev 0.00001072 s 0.000014848 s 0.72
GenDot / IPartOpt / cuda / PreRev 0.000010752 s 0.000009984 s 1.08
GenDot / IPartOpt / cuda / PostRev 0.000010912 s 0.000010176 s 1.07
GenDot / IPartOpt / cuda / BothRev 0.000010943 s 0.000009568 s 1.14
GenDot / DefOpt / cuda / PreRev 0.000011136 s 0.000010751 s 1.04
GenDot / DefOpt / cuda / PostRev 0.000010687 s 0.000010016 s 1.07
GenDot / DefOpt / cuda / BothRev 0.000010656 s 0.000010336 s 1.03
GenDot / IDefOpt / cuda / PreRev 0.000010752 s 0.000010368 s 1.04
GenDot / IDefOpt / cuda / PostRev 0.000010816 s 0.000009984 s 1.08
GenDot / IDefOpt / cuda / BothRev 0.000011232 s 0.000009888 s 1.14
GenDot / JaXPipe / tpu / Primal 9.291e-7 s 9.29575e-7 s 1.00
GenDot / Jax / tpu / Primal 9.259e-7 s 9.25225e-7 s 1.00
GenDot / HLOOpt / tpu / Primal 0.0000015812249999999998 s 0.000001577925 s 1.00
GenDot / PartOpt / tpu / Primal 9.257e-7 s 9.25675e-7 s 1.00
GenDot / IPartOpt / tpu / Primal 9.301e-7 s 9.30175e-7 s 1.00
GenDot / DefOpt / tpu / Primal 0.0000014859749999999998 s 0.00000149485 s 0.99
GenDot / IDefOpt / tpu / Primal 0.0000015697749999999995 s 0.000001579125 s 0.99
GenDot / JaXPipe / tpu / Forward 0.00000316245 s 0.0000031612000000000003 s 1.00
GenDot / Jax / tpu / Forward 0.0000023210000000000003 s 0.000002318925 s 1.00
GenDot / HLOOpt / tpu / Forward 0.000003105925 s 0.0000031131 s 1.00
GenDot / PartOpt / tpu / Forward 0.0000032052500000000004 s 0.00000322215 s 0.99
GenDot / IPartOpt / tpu / Forward 0.000003109725 s 0.000003119 s 1.00
GenDot / DefOpt / tpu / Forward 0.0000032128 s 0.0000032191750000000004 s 1.00
GenDot / IDefOpt / tpu / Forward 0.00000311475 s 0.0000031159 s 1.00
GenDot / JaXPipe / tpu / PreRev 0.00000295635 s 0.000002954675 s 1.00
GenDot / JaXPipe / tpu / PostRev 0.000002408225 s 0.00000241495 s 1.00
GenDot / JaXPipe / tpu / BothRev 0.00000296185 s 0.000002964825 s 1.00
GenDot / Jax / tpu / BothRev 0.0000023982 s 0.000002414275 s 0.99
GenDot / HLOOpt / tpu / PreRev 0.000002959325 s 0.00000296385 s 1.00
GenDot / HLOOpt / tpu / PostRev 0.0000029333 s 0.00000293875 s 1.00
GenDot / HLOOpt / tpu / BothRev 0.000002955425 s 0.000002956875 s 1.00
GenDot / PartOpt / tpu / PreRev 0.000002924075 s 0.000002932425 s 1.00
GenDot / PartOpt / tpu / PostRev 0.0000024007 s 0.00000239235 s 1.00
GenDot / PartOpt / tpu / BothRev 0.0000029353 s 0.0000029372 s 1.00
GenDot / IPartOpt / tpu / PreRev 0.0000029450749999999995 s 0.000002964875 s 0.99
GenDot / IPartOpt / tpu / PostRev 0.000002413475 s 0.000002418575 s 1.00
GenDot / IPartOpt / tpu / BothRev 0.00000295485 s 0.00000296245 s 1.00
GenDot / DefOpt / tpu / PreRev 0.000002932525 s 0.000002938775 s 1.00
GenDot / DefOpt / tpu / PostRev 0.000002954275 s 0.0000029676750000000003 s 1.00
GenDot / DefOpt / tpu / BothRev 0.000002930625 s 0.00000292875 s 1.00
GenDot / IDefOpt / tpu / PreRev 0.000002959625 s 0.000002963075 s 1.00
GenDot / IDefOpt / tpu / PostRev 0.0000029325000000000003 s 0.0000029291 s 1.00
GenDot / IDefOpt / tpu / BothRev 0.0000029544 s 0.0000029657499999999995 s 1.00
GenDot / JaXPipe / cpu / Primal 0.000014477 s 0.000006996919955781777 s 2.07
GenDot / Jax / cpu / Primal 0.000014979 s 0.000007466960050805938 s 2.01
GenDot / HLOOpt / cpu / Primal 0.000014163 s 0.000007573020011477638 s 1.87
GenDot / PartOpt / cpu / Primal 0.000015356 s 0.00000727000000551925 s 2.11
GenDot / IPartOpt / cpu / Primal 0.00001433 s 0.000007262939989232109 s 1.97
GenDot / DefOpt / cpu / Primal 0.000013984 s 0.000007189280004240572 s 1.95
GenDot / IDefOpt / cpu / Primal 0.000014125 s 0.000006937339976502699 s 2.04
GenDot / JaXPipe / cpu / Forward 0.000019415000000000003 s 0.00001109208003981621 s 1.75
GenDot / Jax / cpu / Forward 0.000020372 s 0.000010694279944800656 s 1.90
GenDot / HLOOpt / cpu / Forward 0.000019286 s 0.000011425719985709292 s 1.69
GenDot / PartOpt / cpu / Forward 0.000019116 s 0.000011010379994331744 s 1.74
GenDot / IPartOpt / cpu / Forward 0.000020224 s 0.000011673500002871151 s 1.73
GenDot / DefOpt / cpu / Forward 0.000019464 s 0.000011217820037927595 s 1.74
GenDot / IDefOpt / cpu / Forward 0.000019344 s 0.000011160699987158296 s 1.73
GenDot / JaXPipe / cpu / PreRev 0.000019178 s 0.000011052339996240334 s 1.74
GenDot / JaXPipe / cpu / PostRev 0.000020646 s 0.000010355519953009206 s 1.99
GenDot / JaXPipe / cpu / BothRev 0.000019972 s 0.000011711480019584996 s 1.71
GenDot / Jax / cpu / BothRev 0.000021202 s 0.000010143320014321944 s 2.09
GenDot / HLOOpt / cpu / PreRev 0.000019346 s 0.00001173204003862338 s 1.65
GenDot / HLOOpt / cpu / PostRev 0.000019608 s 0.000013266199948702706 s 1.48
GenDot / HLOOpt / cpu / BothRev 0.000020313 s 0.00001131287996940955 s 1.80
GenDot / PartOpt / cpu / PreRev 0.000019424 s 0.00001143761996900139 s 1.70
GenDot / PartOpt / cpu / PostRev 0.000021865 s 0.000010051760045826088 s 2.18
GenDot / PartOpt / cpu / BothRev 0.000019973 s 0.00001174368006104487 s 1.70
GenDot / IPartOpt / cpu / PreRev 0.000019278 s 0.000011260359970037826 s 1.71
GenDot / IPartOpt / cpu / PostRev 0.000020889000000000003 s 0.000009902019992296118 s 2.11
GenDot / IPartOpt / cpu / BothRev 0.000019635 s 0.000011261819963692687 s 1.74
GenDot / DefOpt / cpu / PreRev 0.000019996 s 0.000011598100036280811 s 1.72
GenDot / DefOpt / cpu / PostRev 0.00001912 s 0.00001124961999266816 s 1.70
GenDot / DefOpt / cpu / BothRev 0.000019724 s 0.000011092819977420732 s 1.78
GenDot / IDefOpt / cpu / PreRev 0.000019754 s 0.000011347700010446716 s 1.74
GenDot / IDefOpt / cpu / PostRev 0.000019382 s 0.000012072240006091306 s 1.61
GenDot / IDefOpt / cpu / BothRev 0.000019292 s 0.00001116171999456128 s 1.73
hlo_ffi / JaXPipe / cpu / Primal 0.000010060140002678963 s 0.000010841679968507378 s 0.93
hlo_ffi / Jax / cpu / Primal 0.00000989259997368208 s 0.000010750079982244642 s 0.92
hlo_ffi / HLOOpt / cpu / Primal 0.00000998584003355063 s 0.000010333600012017995 s 0.97
hlo_ffi / PartOpt / cpu / Primal 0.00000993284003016015 s 0.000010718239973357412 s 0.93
hlo_ffi / IPartOpt / cpu / Primal 0.000010355339991292567 s 0.00001055310000992904 s 0.98
hlo_ffi / DefOpt / cpu / Primal 0.000009707099970910348 s 0.00000993140000900894 s 0.98
hlo_ffi / IDefOpt / cpu / Primal 0.000009776659971976187 s 0.000010487620002095355 s 0.93
hlo_ffi / JaXPipe / cpu / Forward 0.000013981300053274028 s 0.000014204600020093496 s 0.98
hlo_ffi / Jax / cpu / Forward 0.000014093259978835704 s 0.0000146346400197217 s 0.96
hlo_ffi / HLOOpt / cpu / Forward 0.000014227840001694858 s 0.000015109380010471795 s 0.94
hlo_ffi / PartOpt / cpu / Forward 0.000014002219977555796 s 0.000014914740049789544 s 0.94
hlo_ffi / IPartOpt / cpu / Forward 0.000014159900028971606 s 0.000014542079998136617 s 0.97
hlo_ffi / DefOpt / cpu / Forward 0.000014177540015225531 s 0.000015062679985931026 s 0.94
hlo_ffi / IDefOpt / cpu / Forward 0.000014157880023049077 s 0.000014916280015313532 s 0.95
hlo_ffi / JaXPipe / cpu / PreRev 0.00001485058002799633 s 0.00001550459998725273 s 0.96
hlo_ffi / JaXPipe / cpu / PostRev 0.000013662260007549776 s 0.000014995539968367666 s 0.91
hlo_ffi / JaXPipe / cpu / BothRev 0.000014252520004447434 s 0.000014779980065213749 s 0.96
hlo_ffi / Jax / cpu / BothRev 0.00001438293998944573 s 0.0000151172800178756 s 0.95
hlo_ffi / HLOOpt / cpu / PreRev 0.00001429645999451168 s 0.000015309580030589133 s 0.93
hlo_ffi / HLOOpt / cpu / PostRev 0.000016008299999157315 s 0.000016761159986344865 s 0.96
hlo_ffi / HLOOpt / cpu / BothRev 0.000013982780037622432 s 0.000014449539994529914 s 0.97
hlo_ffi / PartOpt / cpu / PreRev 0.000013792799991279026 s 0.000015427959988301153 s 0.89
hlo_ffi / PartOpt / cpu / PostRev 0.000013626460004161344 s 0.000014358119988173713 s 0.95
hlo_ffi / PartOpt / cpu / BothRev 0.000013890180007365416 s 0.000014243580008042045 s 0.98
hlo_ffi / IPartOpt / cpu / PreRev 0.00001404324003488 s 0.000015118039955268616 s 0.93
hlo_ffi / IPartOpt / cpu / PostRev 0.000013435640021270956 s 0.000014513700016323128 s 0.93
hlo_ffi / IPartOpt / cpu / BothRev 0.000013758360018982783 s 0.000014463360030276817 s 0.95
hlo_ffi / DefOpt / cpu / PreRev 0.000014044300005480182 s 0.000014994799976193462 s 0.94
hlo_ffi / DefOpt / cpu / PostRev 0.000014139479999357718 s 0.000014343659986479906 s 0.99
hlo_ffi / DefOpt / cpu / BothRev 0.000013838040013069986 s 0.000014086640021560015 s 0.98
hlo_ffi / IDefOpt / cpu / PreRev 0.00001445914003852522 s 0.000015176240003711427 s 0.95
hlo_ffi / IDefOpt / cpu / PostRev 0.000013767700029347906 s 0.000014299419990493336 s 0.96
hlo_ffi / IDefOpt / cpu / BothRev 0.00001395327999489382 s 0.000014242560055208742 s 0.98
hlo_ffi / JaXPipe / cuda / Primal 0.0000023670000000000004 s 0.000001983 s 1.19
hlo_ffi / Jax / cuda / Primal 0.000002399 s 0.000001952 s 1.23
hlo_ffi / HLOOpt / cuda / Primal 0.000002368 s 0.000001951 s 1.21
hlo_ffi / PartOpt / cuda / Primal 0.000002368 s 0.000001952 s 1.21
hlo_ffi / IPartOpt / cuda / Primal 0.000002368 s 0.000001984 s 1.19
hlo_ffi / DefOpt / cuda / Primal 0.000002399 s 0.000001983 s 1.21
hlo_ffi / IDefOpt / cuda / Primal 0.000002368 s 0.000001983 s 1.19
hlo_ffi / JaXPipe / cuda / Forward 0.000002464 s 0.000002048 s 1.20
hlo_ffi / Jax / cuda / Forward 0.000002463 s 0.000002047 s 1.20
hlo_ffi / HLOOpt / cuda / Forward 0.000002463 s 0.000002047 s 1.20
hlo_ffi / PartOpt / cuda / Forward 0.000002464 s 0.000002047 s 1.20
hlo_ffi / IPartOpt / cuda / Forward 0.000002463 s 0.000002047 s 1.20
hlo_ffi / DefOpt / cuda / Forward 0.000002463 s 0.000002048 s 1.20
hlo_ffi / IDefOpt / cuda / Forward 0.000002464 s 0.000002047 s 1.20
hlo_ffi / JaXPipe / cuda / PreRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / JaXPipe / cuda / PostRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / JaXPipe / cuda / BothRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / Jax / cuda / BothRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / HLOOpt / cuda / PreRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / HLOOpt / cuda / PostRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / HLOOpt / cuda / BothRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / PartOpt / cuda / PreRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / PartOpt / cuda / PostRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / PartOpt / cuda / BothRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / IPartOpt / cuda / PreRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / IPartOpt / cuda / PostRev 0.000002464 s 0.000002047 s 1.20
hlo_ffi / IPartOpt / cuda / BothRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / DefOpt / cuda / PreRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / DefOpt / cuda / PostRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / DefOpt / cuda / BothRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / IDefOpt / cuda / PreRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / IDefOpt / cuda / PostRev 0.000002464 s 0.000002047 s 1.20
hlo_ffi / IDefOpt / cuda / BothRev 0.000002463 s 0.000002047 s 1.20
hlo_ffi / JaXPipe / tpu / Primal 9.32025e-7 s 9.225e-7 s 1.01
hlo_ffi / Jax / tpu / Primal 9.5485e-7 s 9.495e-7 s 1.01
hlo_ffi / HLOOpt / tpu / Primal 9.07e-7 s 9.03125e-7 s 1.00
hlo_ffi / PartOpt / tpu / Primal 9.5535e-7 s 9.57925e-7 s 1.00
hlo_ffi / IPartOpt / tpu / Primal 9.0845e-7 s 9.03575e-7 s 1.01
hlo_ffi / DefOpt / tpu / Primal 9.54975e-7 s 9.537e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Primal 9.081e-7 s 8.97925e-7 s 1.01
hlo_ffi / JaXPipe / tpu / Forward 9.4905e-7 s 9.486e-7 s 1.00
hlo_ffi / Jax / tpu / Forward 9.8135e-7 s 9.813e-7 s 1.00
hlo_ffi / HLOOpt / tpu / Forward 9.74125e-7 s 9.736e-7 s 1.00
hlo_ffi / PartOpt / tpu / Forward 9.34075e-7 s 9.33875e-7 s 1.00
hlo_ffi / IPartOpt / tpu / Forward 9.74175e-7 s 9.73325e-7 s 1.00
hlo_ffi / DefOpt / tpu / Forward 9.33775e-7 s 9.337e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Forward 9.743499999999998e-7 s 9.74125e-7 s 1.00
hlo_ffi / JaXPipe / tpu / PreRev 9.37875e-7 s 9.317e-7 s 1.01
hlo_ffi / JaXPipe / tpu / PostRev 9.65425e-7 s 9.64725e-7 s 1.00
hlo_ffi / JaXPipe / tpu / BothRev 9.6165e-7 s 9.619e-7 s 1.00
hlo_ffi / Jax / tpu / BothRev 9.652e-7 s 9.6445e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PreRev 9.623e-7 s 9.61625e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PostRev 9.654e-7 s 9.650749999999998e-7 s 1.00
hlo_ffi / HLOOpt / tpu / BothRev 9.62375e-7 s 9.615e-7 s 1.00
hlo_ffi / PartOpt / tpu / PreRev 9.65525e-7 s 9.6405e-7 s 1.00
hlo_ffi / PartOpt / tpu / PostRev 9.627e-7 s 9.61925e-7 s 1.00
hlo_ffi / PartOpt / tpu / BothRev 9.64975e-7 s 9.6445e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PreRev 9.625250000000002e-7 s 9.61975e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PostRev 9.653e-7 s 9.65e-7 s 1.00
hlo_ffi / IPartOpt / tpu / BothRev 9.619e-7 s 9.62325e-7 s 1.00
hlo_ffi / DefOpt / tpu / PreRev 9.65e-7 s 9.64275e-7 s 1.00
hlo_ffi / DefOpt / tpu / PostRev 9.624e-7 s 9.62e-7 s 1.00
hlo_ffi / DefOpt / tpu / BothRev 9.65575e-7 s 9.63975e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PreRev 9.618500000000002e-7 s 9.62075e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PostRev 9.64875e-7 s 9.64625e-7 s 1.00
hlo_ffi / IDefOpt / tpu / BothRev 9.61775e-7 s 9.61825e-7 s 1.00
hlo_ffi / JaXPipe / cpu / Primal 0.000017796 s 0.000010841679968507378 s 1.64
hlo_ffi / Jax / cpu / Primal 0.00001755 s 0.000010750079982244642 s 1.63
hlo_ffi / HLOOpt / cpu / Primal 0.000016971 s 0.000010333600012017995 s 1.64
hlo_ffi / PartOpt / cpu / Primal 0.000017825 s 0.000010718239973357412 s 1.66
hlo_ffi / IPartOpt / cpu / Primal 0.000017804 s 0.00001055310000992904 s 1.69
hlo_ffi / DefOpt / cpu / Primal 0.000017907 s 0.00000993140000900894 s 1.80
hlo_ffi / IDefOpt / cpu / Primal 0.000017548999999999997 s 0.000010487620002095355 s 1.67
hlo_ffi / JaXPipe / cpu / Forward 0.000025291 s 0.000014204600020093496 s 1.78
hlo_ffi / Jax / cpu / Forward 0.00002472 s 0.0000146346400197217 s 1.69
hlo_ffi / HLOOpt / cpu / Forward 0.000024593 s 0.000015109380010471795 s 1.63
hlo_ffi / PartOpt / cpu / Forward 0.000024231 s 0.000014914740049789544 s 1.62
hlo_ffi / IPartOpt / cpu / Forward 0.000024613 s 0.000014542079998136617 s 1.69
hlo_ffi / DefOpt / cpu / Forward 0.000026154 s 0.000015062679985931026 s 1.74
hlo_ffi / IDefOpt / cpu / Forward 0.000024316 s 0.000014916280015313532 s 1.63
hlo_ffi / JaXPipe / cpu / PreRev 0.00002409 s 0.00001550459998725273 s 1.55
hlo_ffi / JaXPipe / cpu / PostRev 0.000024331 s 0.000014995539968367666 s 1.62
hlo_ffi / JaXPipe / cpu / BothRev 0.000023948 s 0.000014779980065213749 s 1.62
hlo_ffi / Jax / cpu / BothRev 0.000024726 s 0.0000151172800178756 s 1.64
hlo_ffi / HLOOpt / cpu / PreRev 0.00002355 s 0.000015309580030589133 s 1.54
hlo_ffi / HLOOpt / cpu / PostRev 0.000024613 s 0.000016761159986344865 s 1.47
hlo_ffi / HLOOpt / cpu / BothRev 0.000025085 s 0.000014449539994529914 s 1.74
hlo_ffi / PartOpt / cpu / PreRev 0.000024285 s 0.000015427959988301153 s 1.57
hlo_ffi / PartOpt / cpu / PostRev 0.000023953 s 0.000014358119988173713 s 1.67
hlo_ffi / PartOpt / cpu / BothRev 0.000024318 s 0.000014243580008042045 s 1.71
hlo_ffi / IPartOpt / cpu / PreRev 0.000024744 s 0.000015118039955268616 s 1.64
hlo_ffi / IPartOpt / cpu / PostRev 0.000024197 s 0.000014513700016323128 s 1.67
hlo_ffi / IPartOpt / cpu / BothRev 0.000024716 s 0.000014463360030276817 s 1.71
hlo_ffi / DefOpt / cpu / PreRev 0.000024504 s 0.000014994799976193462 s 1.63
hlo_ffi / DefOpt / cpu / PostRev 0.000024716 s 0.000014343659986479906 s 1.72
hlo_ffi / DefOpt / cpu / BothRev 0.000025029 s 0.000014086640021560015 s 1.78
hlo_ffi / IDefOpt / cpu / PreRev 0.000024512 s 0.000015176240003711427 s 1.62
hlo_ffi / IDefOpt / cpu / PostRev 0.000024334 s 0.000014299419990493336 s 1.70
hlo_ffi / IDefOpt / cpu / BothRev 0.00002422 s 0.000014242560055208742 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.0009186375999888 s 0.0009240977999979 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 / cpu / Primal 0.0009215183999003 s 0.0009339908000583 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 / cpu / Primal 0.0009583683999153 s 0.0009847250000348 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 / PartOpt / cpu / Primal 0.0009158448000562 s 0.0009646690000408 s 0.95
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.0009349971998744 s 0.0009248692001165 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 / Primal 0.0009527468000669 s 0.0010179868000705 s 0.94
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.0009667225999692 s 0.0010224614000435 s 0.95
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.002137538199986 s 0.0023158776000855 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 / Jax / cpu / Forward 0.0024268912001389 s 0.0024215126000854 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 / Forward 0.0021988163999594 s 0.0024494521998349 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 / PartOpt / cpu / Forward 0.0021820317999299 s 0.0023676605998844 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 / IPartOpt / cpu / Forward 0.0021992990001308 s 0.0023951875999046 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 / DefOpt / cpu / Forward 0.0021680532000573 s 0.0023946858000272 s 0.91
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.0021736877999501 s 0.0024209647998759 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 / JaXPipe / cpu / PreRev 0.0052462801999354 s 0.006682462200024 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 / PostRev 0.0064046144000712 s 0.0063939480000044 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 / BothRev 0.0056180652000875 s 0.0062240333999397 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.0064315644000998 s 0.006168826800058 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 / HLOOpt / cpu / PreRev 0.005430910600171 s 0.0058642033999603 s 0.93
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.0068446151999523 s 0.0065649374000713 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 / HLOOpt / cpu / BothRev 0.0046433861998593 s 0.0062589523998212 s 0.74
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.0051953873999082 s 0.0055141665999144 s 0.94
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.0061607350000485 s 0.006519494800159 s 0.94
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.0055189122000228 s 0.0060325897999973 s 0.91
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.0053517284000008 s 0.0058518562001154 s 0.91
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.0055126527999163 s 0.0066064911999092 s 0.83
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.0052459774000453 s 0.0049009999999725 s 1.07
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.0052196981999259 s 0.0068223805998968 s 0.77
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.0049885803999131 s 0.0042839250001634 s 1.16
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.0052707517999806 s 0.0049773534000451 s 1.06
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.0052950569999666 s 0.0050564379999741 s 1.05
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.0052874332000101 s 0.0036666862000856 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 / IDefOpt / cpu / BothRev 0.0052458469999692 s 0.0049634575999334 s 1.06
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.000299325 s 0.000282525 s 1.06
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.000300638 s 0.000282142 s 1.07
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.000305565 s 0.000288765 s 1.06
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.000299837 s 0.000281629 s 1.06
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.000299165 s 0.000282238 s 1.06
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.000306653 s 0.000289949 s 1.06
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.000305757 s 0.000288957 s 1.06
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.000584218 s 0.000560476 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 / Jax / cuda / Forward 0.000568155 s 0.000540347 s 1.05
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.000584379 s 0.000560635 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 / PartOpt / cuda / Forward 0.000583738 s 0.0005582359999999 s 1.05
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.00058457 s 0.0005590349999999 s 1.05
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.000584026 s 0.000558843 s 1.05
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.000584218 s 0.0005579469999999 s 1.05
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.001056117 s 0.001036343 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 / cuda / PostRev 0.00101359 s 0.000989623 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 / cuda / BothRev 0.001054678 s 0.001025015 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 / cuda / BothRev 0.001007383 s 0.000996023 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 / cuda / PreRev 0.001042262 s 0.001019096 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 / PostRev 0.001064982 s 0.001040856 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 / BothRev 0.001043062 s 0.001012472 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 / cuda / PreRev 0.001053014 s 0.001032887 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 / cuda / PostRev 0.0010015909999999 s 0.000979896 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 / cuda / BothRev 0.001052086 s 0.001029496 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 / IPartOpt / cuda / PreRev 0.001053526 s 0.001029495 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 / IPartOpt / cuda / PostRev 0.001004055 s 0.000981815 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 / IPartOpt / cuda / BothRev 0.0010520549999999 s 0.0010351269999999 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 / cuda / PreRev 0.001055894 s 0.001031063 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 / cuda / PostRev 0.000989591 s 0.0009678 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 / cuda / BothRev 0.00105743 s 0.001030904 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 / IDefOpt / cuda / PreRev 0.001056886 s 0.001029335 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 / IDefOpt / cuda / PostRev 0.001057846 s 0.001028791 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 / IDefOpt / cuda / BothRev 0.001056438 s 0.001029911 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 / JaXPipe / tpu / Primal 0.00012363575 s 0.0001236612499999 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 / Primal 0.00012666575 s 0.00012691525 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 / Primal 0.00015263225 s 0.00015220075 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 / Primal 0.00013397125 s 0.00013378675 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 / Primal 0.0001311605 s 0.000131553 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 / Primal 0.0001477715 s 0.00014771625 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 / Primal 0.00015075225 s 0.00015105825 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 / Forward 0.0002121329999999 s 0.00021229 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.000261353 s 0.00026106525 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.00021183025 s 0.0002123467499999 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 / Forward 0.0002181555 s 0.000218429 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 / Forward 0.00021181675 s 0.00021238525 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 / Forward 0.0002182717499999 s 0.0002185895 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 / Forward 0.00021203775 s 0.00021234975 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 / PreRev 0.00035660875 s 0.00035685875 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.0002585224999999 s 0.00025662975 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 / BothRev 0.00035646575 s 0.00035696275 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.00026016325 s 0.0002577185 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 / PreRev 0.000356696 s 0.00035700625 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.00029211775 s 0.0002913595 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.000356719 s 0.0003571645 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.0003586337499999 s 0.00035785075 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.00027194225 s 0.0002721935 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 / BothRev 0.00035854425 s 0.000358094 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.0003564265 s 0.0003567829999999 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.0002758285 s 0.00027296725 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 / BothRev 0.00035645625 s 0.00035733875 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.0003597675 s 0.00035909675 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.0002838987499999 s 0.0002837755 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.0003596 s 0.00035903975 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.00035800575 s 0.00035801775 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.000303014 s 0.00030154675 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.000358169 s 0.00035856175 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.002291955 s 0.0009240977999979 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 / Jax / cpu / Primal 0.002391911 s 0.0009339908000583 s 2.56
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.002278025 s 0.0009847250000348 s 2.31
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.002315106 s 0.0009646690000408 s 2.40
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.002384122 s 0.0009248692001165 s 2.58
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.002496941 s 0.0010179868000705 s 2.45
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.00301772 s 0.0010224614000435 s 2.95
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.005348498 s 0.0023158776000855 s 2.31
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.005738817 s 0.0024215126000854 s 2.37
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.005747464 s 0.0024494521998349 s 2.35
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.005550542 s 0.0023676605998844 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 / Forward 0.005802517 s 0.0023951875999046 s 2.42
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.005622282 s 0.0023946858000272 s 2.35
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.005957149 s 0.0024209647998759 s 2.46
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.009068559 s 0.006682462200024 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 / JaXPipe / cpu / PostRev 0.008525033 s 0.0063939480000044 s 1.33
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.00896358 s 0.0062240333999397 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 / Jax / cpu / BothRev 0.0091387669999999 s 0.006168826800058 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 / HLOOpt / cpu / PreRev 0.009605341 s 0.0058642033999603 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 / HLOOpt / cpu / PostRev 0.009309608 s 0.0065649374000713 s 1.42
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.0097479949999999 s 0.0062589523998212 s 1.56
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.009526736 s 0.0055141665999144 s 1.73
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.009463763 s 0.006519494800159 s 1.45
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.008970578 s 0.0060325897999973 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 / IPartOpt / cpu / PreRev 0.009469891 s 0.0058518562001154 s 1.62
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.00945354 s 0.0066064911999092 s 1.43
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.008566504 s 0.0049009999999725 s 1.75
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.009624509 s 0.0068223805998968 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 / DefOpt / cpu / PostRev 0.0083295949999999 s 0.0042839250001634 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 / DefOpt / cpu / BothRev 0.00873563 s 0.0049773534000451 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 / IDefOpt / cpu / PreRev 0.008656967 s 0.0050564379999741 s 1.71
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.009007899 s 0.0036666862000856 s 2.46
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.008143839 s 0.0049634575999334 s 1.64
scatter_sum / JaXPipe / cpu / Primal 0.000008843880059430377 s 0.000007927520009616273 s 1.12
scatter_sum / Jax / cpu / Primal 0.000008500980002281721 s 0.00000742767999327043 s 1.14
scatter_sum / HLOOpt / cpu / Primal 0.000008701579999979004 s 0.000008308480018968112 s 1.05
scatter_sum / PartOpt / cpu / Primal 0.00000868269998136384 s 0.000007928900022307061 s 1.10
scatter_sum / IPartOpt / cpu / Primal 0.000008923260011215462 s 0.000007562780001535429 s 1.18
scatter_sum / DefOpt / cpu / Primal 0.000008373480013688095 s 0.00000764010001148563 s 1.10
scatter_sum / IDefOpt / cpu / Primal 0.000008563639985368354 s 0.000007460480064764852 s 1.15
scatter_sum / JaXPipe / cpu / Forward 0.000013607460023195017 s 0.000011345139937475325 s 1.20
scatter_sum / Jax / cpu / Forward 0.00001434909997442446 s 0.000011204619995623944 s 1.28
scatter_sum / HLOOpt / cpu / Forward 0.000014884379952491145 s 0.00001257889996850281 s 1.18
scatter_sum / PartOpt / cpu / Forward 0.000013651439985551406 s 0.00001152361997810658 s 1.18
scatter_sum / IPartOpt / cpu / Forward 0.000014551259973814012 s 0.000014340680008899652 s 1.01
scatter_sum / DefOpt / cpu / Forward 0.000014451719980570488 s 0.000011413759984861828 s 1.27
scatter_sum / IDefOpt / cpu / Forward 0.000014404939975065645 s 0.000011250819998167572 s 1.28
scatter_sum / JaXPipe / cpu / PreRev 0.000014194780005709616 s 0.000011935600005017476 s 1.19
scatter_sum / JaXPipe / cpu / PostRev 0.000013157999974282575 s 0.000011985679984718445 s 1.10
scatter_sum / JaXPipe / cpu / BothRev 0.00001379985999847122 s 0.000012393960014378536 s 1.11
scatter_sum / Jax / cpu / BothRev 0.000013558219943661245 s 0.000011429759997554356 s 1.19
scatter_sum / HLOOpt / cpu / PreRev 0.000012971120013389736 s 0.000012093140012439108 s 1.07
scatter_sum / HLOOpt / cpu / PostRev 0.00001583037999807857 s 0.000014394499994523355 s 1.10
scatter_sum / HLOOpt / cpu / BothRev 0.000013417139962257352 s 0.000011507600011100294 s 1.17
scatter_sum / PartOpt / cpu / PreRev 0.000013496500005203415 s 0.00001178795998384885 s 1.14
scatter_sum / PartOpt / cpu / PostRev 0.000013723280035264909 s 0.00001198563998514146 s 1.14
scatter_sum / PartOpt / cpu / BothRev 0.00001433055996130861 s 0.000012446580021787668 s 1.15
scatter_sum / IPartOpt / cpu / PreRev 0.00001317157999437768 s 0.000011264420027146117 s 1.17
scatter_sum / IPartOpt / cpu / PostRev 0.000012723239988190472 s 0.000011778660000345554 s 1.08
scatter_sum / IPartOpt / cpu / BothRev 0.000012858120007877006 s 0.000011508220040923334 s 1.12
scatter_sum / DefOpt / cpu / PreRev 0.000013342600013857008 s 0.000011381679996702587 s 1.17
scatter_sum / DefOpt / cpu / PostRev 0.0000129961800394085 s 0.0000118991999715945 s 1.09
scatter_sum / DefOpt / cpu / BothRev 0.000013398019973465123 s 0.00001180374005343765 s 1.14
scatter_sum / IDefOpt / cpu / PreRev 0.00001326251998762018 s 0.00001141565998295846 s 1.16
scatter_sum / IDefOpt / cpu / PostRev 0.000013449720026983414 s 0.000011669380010062014 s 1.15
scatter_sum / IDefOpt / cpu / BothRev 0.000014036879965715343 s 0.000012216939976497088 s 1.15
scatter_sum / JaXPipe / cuda / Primal 0.00001024 s 0.000010496 s 0.98
scatter_sum / Jax / cuda / Primal 0.000011871 s 0.00000992 s 1.20
scatter_sum / HLOOpt / cuda / Primal 0.000010496 s 0.000010303 s 1.02
scatter_sum / PartOpt / cuda / Primal 0.000010464 s 0.000010144 s 1.03
scatter_sum / IPartOpt / cuda / Primal 0.000011359 s 0.000010304 s 1.10
scatter_sum / DefOpt / cuda / Primal 0.000011808 s 0.00001008 s 1.17
scatter_sum / IDefOpt / cuda / Primal 0.00001168 s 0.000010495 s 1.11
scatter_sum / JaXPipe / cuda / Forward 0.000017408 s 0.00001728 s 1.01
scatter_sum / Jax / cuda / Forward 0.000017312 s 0.000017247999999999998 s 1.00
scatter_sum / HLOOpt / cuda / Forward 0.000019584 s 0.000016927999999999998 s 1.16
scatter_sum / PartOpt / cuda / Forward 0.00001984 s 0.000016576000000000002 s 1.20
scatter_sum / IPartOpt / cuda / Forward 0.000018912 s 0.000016864 s 1.12
scatter_sum / DefOpt / cuda / Forward 0.000019744000000000003 s 0.000017056 s 1.16
scatter_sum / IDefOpt / cuda / Forward 0.000020384 s 0.00001632 s 1.25
scatter_sum / JaXPipe / cuda / PreRev 0.000018528 s 0.000017024 s 1.09
scatter_sum / JaXPipe / cuda / PostRev 0.000017408 s 0.000018431 s 0.94
scatter_sum / JaXPipe / cuda / BothRev 0.00001728 s 0.000016704 s 1.03
scatter_sum / Jax / cuda / BothRev 0.00001712 s 0.000017088 s 1.00
scatter_sum / HLOOpt / cuda / PreRev 0.000017823 s 0.000016672 s 1.07
scatter_sum / HLOOpt / cuda / PostRev 0.000017408 s 0.000016448000000000002 s 1.06
scatter_sum / HLOOpt / cuda / BothRev 0.000017024 s 0.0000168 s 1.01
scatter_sum / PartOpt / cuda / PreRev 0.000017247999999999998 s 0.000016927999999999998 s 1.02
scatter_sum / PartOpt / cuda / PostRev 0.000018112 s 0.000016703 s 1.08
scatter_sum / PartOpt / cuda / BothRev 0.000017408 s 0.00001728 s 1.01
scatter_sum / IPartOpt / cuda / PreRev 0.000018304 s 0.0000168 s 1.09
scatter_sum / IPartOpt / cuda / PostRev 0.000017279 s 0.000017024 s 1.01
scatter_sum / IPartOpt / cuda / BothRev 0.00001696 s 0.000016736 s 1.01
scatter_sum / DefOpt / cuda / PreRev 0.00001728 s 0.0000176 s 0.98
scatter_sum / DefOpt / cuda / PostRev 0.000018912 s 0.000016383999999999998 s 1.15
scatter_sum / DefOpt / cuda / BothRev 0.000017184 s 0.000016542999999999997 s 1.04
scatter_sum / IDefOpt / cuda / PreRev 0.000017344 s 0.000016736 s 1.04
scatter_sum / IDefOpt / cuda / PostRev 0.000017152 s 0.0000168 s 1.02
scatter_sum / IDefOpt / cuda / BothRev 0.000017408 s 0.00001696 s 1.03
scatter_sum / JaXPipe / tpu / Primal 0.00000135125 s 0.000001343475 s 1.01
scatter_sum / Jax / tpu / Primal 0.0000014045500000000002 s 0.0000014044499999999995 s 1.00
scatter_sum / HLOOpt / tpu / Primal 0.000001350625 s 0.000001343175 s 1.01
scatter_sum / PartOpt / tpu / Primal 0.00000140485 s 0.000001404525 s 1.00
scatter_sum / IPartOpt / tpu / Primal 0.00000135135 s 0.0000013431 s 1.01
scatter_sum / DefOpt / tpu / Primal 0.000001405525 s 0.0000014045 s 1.00
scatter_sum / IDefOpt / tpu / Primal 0.0000013512 s 0.00000134275 s 1.01
scatter_sum / JaXPipe / tpu / Forward 0.0000027076750000000003 s 0.000002702225 s 1.00
scatter_sum / Jax / tpu / Forward 0.0000027259 s 0.000002723475 s 1.00
scatter_sum / HLOOpt / tpu / Forward 0.00000270555 s 0.00000270605 s 1.00
scatter_sum / PartOpt / tpu / Forward 0.000002698175 s 0.000002685575 s 1.00
scatter_sum / IPartOpt / tpu / Forward 0.00000270585 s 0.000002701075 s 1.00
scatter_sum / DefOpt / tpu / Forward 0.0000026989999999999995 s 0.000002689975 s 1.00
scatter_sum / IDefOpt / tpu / Forward 0.000002702975 s 0.000002707075 s 1.00
scatter_sum / JaXPipe / tpu / PreRev 0.000002695075 s 0.00000268835 s 1.00
scatter_sum / JaXPipe / tpu / PostRev 0.000002689675 s 0.000002688925 s 1.00
scatter_sum / JaXPipe / tpu / BothRev 0.0000027057 s 0.00000270215 s 1.00
scatter_sum / Jax / tpu / BothRev 0.000002745225 s 0.0000027393 s 1.00
scatter_sum / HLOOpt / tpu / PreRev 0.00000271045 s 0.000002698875 s 1.00
scatter_sum / HLOOpt / tpu / PostRev 0.000002740875 s 0.000002741625 s 1.00
scatter_sum / HLOOpt / tpu / BothRev 0.000002711025 s 0.000002709775 s 1.00
scatter_sum / PartOpt / tpu / PreRev 0.000002755125 s 0.0000027427000000000004 s 1.00
scatter_sum / PartOpt / tpu / PostRev 0.000002703275 s 0.0000026978249999999994 s 1.00
scatter_sum / PartOpt / tpu / BothRev 0.000002740025 s 0.000002742175 s 1.00
scatter_sum / IPartOpt / tpu / PreRev 0.000002703925 s 0.00000270555 s 1.00
scatter_sum / IPartOpt / tpu / PostRev 0.000002740125 s 0.000002738875 s 1.00
scatter_sum / IPartOpt / tpu / BothRev 0.0000027066 s 0.000002701575 s 1.00
scatter_sum / DefOpt / tpu / PreRev 0.0000027443 s 0.0000027513 s 1.00
scatter_sum / DefOpt / tpu / PostRev 0.0000027014 s 0.000002702125 s 1.00
scatter_sum / DefOpt / tpu / BothRev 0.000002750525 s 0.0000027480500000000003 s 1.00
scatter_sum / IDefOpt / tpu / PreRev 0.00000270985 s 0.00000270145 s 1.00
scatter_sum / IDefOpt / tpu / PostRev 0.000002740725 s 0.00000274525 s 1.00
scatter_sum / IDefOpt / tpu / BothRev 0.000002709125 s 0.000002697975 s 1.00
scatter_sum / JaXPipe / cpu / Primal 0.000015858 s 0.000007927520009616273 s 2.00
scatter_sum / Jax / cpu / Primal 0.000015527 s 0.00000742767999327043 s 2.09
scatter_sum / HLOOpt / cpu / Primal 0.000015946999999999998 s 0.000008308480018968112 s 1.92
scatter_sum / PartOpt / cpu / Primal 0.000015997 s 0.000007928900022307061 s 2.02
scatter_sum / IPartOpt / cpu / Primal 0.000015403 s 0.000007562780001535429 s 2.04
scatter_sum / DefOpt / cpu / Primal 0.000015807000000000003 s 0.00000764010001148563 s 2.07
scatter_sum / IDefOpt / cpu / Primal 0.000015481 s 0.000007460480064764852 s 2.08
scatter_sum / JaXPipe / cpu / Forward 0.000023234 s 0.000011345139937475325 s 2.05
scatter_sum / Jax / cpu / Forward 0.000023484 s 0.000011204619995623944 s 2.10
scatter_sum / HLOOpt / cpu / Forward 0.000022595 s 0.00001257889996850281 s 1.80
scatter_sum / PartOpt / cpu / Forward 0.000022929 s 0.00001152361997810658 s 1.99
scatter_sum / IPartOpt / cpu / Forward 0.00002439 s 0.000014340680008899652 s 1.70
scatter_sum / DefOpt / cpu / Forward 0.000023038 s 0.000011413759984861828 s 2.02
scatter_sum / IDefOpt / cpu / Forward 0.0000225 s 0.000011250819998167572 s 2.00
scatter_sum / JaXPipe / cpu / PreRev 0.00002376 s 0.000011935600005017476 s 1.99
scatter_sum / JaXPipe / cpu / PostRev 0.000023381 s 0.000011985679984718445 s 1.95
scatter_sum / JaXPipe / cpu / BothRev 0.000022636 s 0.000012393960014378536 s 1.83
scatter_sum / Jax / cpu / BothRev 0.000022375 s 0.000011429759997554356 s 1.96
scatter_sum / HLOOpt / cpu / PreRev 0.000022908 s 0.000012093140012439108 s 1.89
scatter_sum / HLOOpt / cpu / PostRev 0.000022918 s 0.000014394499994523355 s 1.59
scatter_sum / HLOOpt / cpu / BothRev 0.000022645 s 0.000011507600011100294 s 1.97
scatter_sum / PartOpt / cpu / PreRev 0.00002261 s 0.00001178795998384885 s 1.92
scatter_sum / PartOpt / cpu / PostRev 0.000023227 s 0.00001198563998514146 s 1.94
scatter_sum / PartOpt / cpu / BothRev 0.000022888 s 0.000012446580021787668 s 1.84
scatter_sum / IPartOpt / cpu / PreRev 0.000022909 s 0.000011264420027146117 s 2.03
scatter_sum / IPartOpt / cpu / PostRev 0.000023389 s 0.000011778660000345554 s 1.99
scatter_sum / IPartOpt / cpu / BothRev 0.000023564 s 0.000011508220040923334 s 2.05
scatter_sum / DefOpt / cpu / PreRev 0.000023279 s 0.000011381679996702587 s 2.05
scatter_sum / DefOpt / cpu / PostRev 0.000023111 s 0.0000118991999715945 s 1.94
scatter_sum / DefOpt / cpu / BothRev 0.000023801 s 0.00001180374005343765 s 2.02
scatter_sum / IDefOpt / cpu / PreRev 0.000024071 s 0.00001141565998295846 s 2.11
scatter_sum / IDefOpt / cpu / PostRev 0.000023031 s 0.000011669380010062014 s 1.97
scatter_sum / IDefOpt / cpu / BothRev 0.000023733 s 0.000012216939976497088 s 1.94
slicing / JaXPipe / cpu / Primal 0.000006886999981361441 s 0.0000065718999940145295 s 1.05
slicing / Jax / cpu / Primal 0.000007202140004665125 s 0.000006047020051482832 s 1.19
slicing / HLOOpt / cpu / Primal 0.00000716288001058274 s 0.000006182960005389759 s 1.16
slicing / PartOpt / cpu / Primal 0.000006891999964864226 s 0.000006284980026975973 s 1.10
slicing / IPartOpt / cpu / Primal 0.000007695439999224618 s 0.0000066760800291376655 s 1.15
slicing / DefOpt / cpu / Primal 0.000006781920046705636 s 0.000006011119994582259 s 1.13
slicing / IDefOpt / cpu / Primal 0.000006458039988501696 s 0.00000617611998677603 s 1.05
slicing / JaXPipe / cpu / Forward 0.000010714520003602956 s 0.00000992725998003152 s 1.08
slicing / Jax / cpu / Forward 0.000010524719973545873 s 0.000009392720012328936 s 1.12
slicing / HLOOpt / cpu / Forward 0.000010868519993891824 s 0.000010251100020468584 s 1.06
slicing / PartOpt / cpu / Forward 0.000010317340002075071 s 0.000009866499995041525 s 1.05
slicing / IPartOpt / cpu / Forward 0.000010905479994107737 s 0.000010019239998655394 s 1.09
slicing / DefOpt / cpu / Forward 0.000010989259963025688 s 0.000009773899992069344 s 1.12
slicing / IDefOpt / cpu / Forward 0.000010210780019406229 s 0.000009806559974094852 s 1.04
slicing / JaXPipe / cpu / PreRev 0.00001127814000028593 s 0.000010159940056837514 s 1.11
slicing / JaXPipe / cpu / PostRev 0.000011219879961572588 s 0.000010119239987034234 s 1.11
slicing / JaXPipe / cpu / BothRev 0.000011716779981725267 s 0.000010195520026172744 s 1.15
slicing / Jax / cpu / BothRev 0.00001102534000892774 s 0.000010721000016928885 s 1.03
slicing / HLOOpt / cpu / PreRev 0.000011662540000543233 s 0.000010868420013139255 s 1.07
slicing / HLOOpt / cpu / PostRev 0.000013378539979385095 s 0.000012423019970810856 s 1.08
slicing / HLOOpt / cpu / BothRev 0.000011141719987790566 s 0.000009850219994405052 s 1.13
slicing / PartOpt / cpu / PreRev 0.00001127222000832262 s 0.00000993288001154724 s 1.13
slicing / PartOpt / cpu / PostRev 0.000011227479999433854 s 0.000010278080017087632 s 1.09
slicing / PartOpt / cpu / BothRev 0.000011293119978290634 s 0.00001018304002172954 s 1.11
slicing / IPartOpt / cpu / PreRev 0.000011266539959251533 s 0.000010028739980043611 s 1.12
slicing / IPartOpt / cpu / PostRev 0.000011409299995648324 s 0.00001013477996821166 s 1.13
slicing / IPartOpt / cpu / BothRev 0.000011180999981661444 s 0.000009986599998228483 s 1.12
slicing / DefOpt / cpu / PreRev 0.000011055820032197516 s 0.000010073839976030283 s 1.10
slicing / DefOpt / cpu / PostRev 0.000011696140009007647 s 0.000010306819967809134 s 1.13
slicing / DefOpt / cpu / BothRev 0.000011523219964146848 s 0.000009928880026564 s 1.16
slicing / IDefOpt / cpu / PreRev 0.00001122604001466243 s 0.000010205140033576756 s 1.10
slicing / IDefOpt / cpu / PostRev 0.000011205980026716132 s 0.00001020074000734894 s 1.10
slicing / IDefOpt / cpu / BothRev 0.000010760999994090525 s 0.000010028259966929908 s 1.07
slicing / JaXPipe / cuda / Primal 0.000002304 s 0.000001887 s 1.22
slicing / Jax / cuda / Primal 0.000002303 s 0.000001888 s 1.22
slicing / HLOOpt / cuda / Primal 0.000002304 s 0.000001887 s 1.22
slicing / PartOpt / cuda / Primal 0.000002304 s 0.000001888 s 1.22
slicing / IPartOpt / cuda / Primal 0.000002304 s 0.000001888 s 1.22
slicing / DefOpt / cuda / Primal 0.000002304 s 0.000001887 s 1.22
slicing / IDefOpt / cuda / Primal 0.000002304 s 0.000001887 s 1.22
slicing / JaXPipe / cuda / Forward 0.000010432 s 0.000010208 s 1.02
slicing / Jax / cuda / Forward 0.00001024 s 0.00001008 s 1.02
slicing / HLOOpt / cuda / Forward 0.000010176 s 0.000010015 s 1.02
slicing / PartOpt / cuda / Forward 0.000010016 s 0.000010176 s 0.98
slicing / IPartOpt / cuda / Forward 0.000010368 s 0.000009664 s 1.07
slicing / DefOpt / cuda / Forward 0.000011103 s 0.000010048 s 1.10
slicing / IDefOpt / cuda / Forward 0.000010304 s 0.000009984 s 1.03
slicing / JaXPipe / cuda / PreRev 0.000010336 s 0.000009695 s 1.07
slicing / JaXPipe / cuda / PostRev 0.000010432 s 0.000009696 s 1.08
slicing / JaXPipe / cuda / BothRev 0.000010463 s 0.000009791 s 1.07
slicing / Jax / cuda / BothRev 0.000011584 s 0.000009536 s 1.21
slicing / HLOOpt / cuda / PreRev 0.000010688 s 0.000009792 s 1.09
slicing / HLOOpt / cuda / PostRev 0.000010496 s 0.000009311 s 1.13
slicing / HLOOpt / cuda / BothRev 0.000010336 s 0.000009824 s 1.05
slicing / PartOpt / cuda / PreRev 0.000010624 s 0.000010208 s 1.04
slicing / PartOpt / cuda / PostRev 0.000010432 s 0.00001008 s 1.03
slicing / PartOpt / cuda / BothRev 0.000010816 s 0.00001008 s 1.07
slicing / IPartOpt / cuda / PreRev 0.000010496 s 0.000010049 s 1.04
slicing / IPartOpt / cuda / PostRev 0.000010303 s 0.00001072 s 0.96
slicing / IPartOpt / cuda / BothRev 0.000010784 s 0.000009505 s 1.13
slicing / DefOpt / cuda / PreRev 0.000010592 s 0.000010016 s 1.06
slicing / DefOpt / cuda / PostRev 0.000010368 s 0.000010208 s 1.02
slicing / DefOpt / cuda / BothRev 0.000010592 s 0.00001024 s 1.03
slicing / IDefOpt / cuda / PreRev 0.000010719 s 0.000010144 s 1.06
slicing / IDefOpt / cuda / PostRev 0.000010176 s 0.00000976 s 1.04
slicing / IDefOpt / cuda / BothRev 0.000010304 s 0.000010111 s 1.02
slicing / JaXPipe / tpu / Primal 0.000001027825 s 0.000001018175 s 1.01
slicing / Jax / tpu / Primal 9.68275e-7 s 9.61625e-7 s 1.01
slicing / HLOOpt / tpu / Primal 0.00000102385 s 0.0000010201000000000002 s 1.00
slicing / PartOpt / tpu / Primal 9.74025e-7 s 9.68875e-7 s 1.01
slicing / IPartOpt / tpu / Primal 0.00000102995 s 0.0000010229 s 1.01
slicing / DefOpt / tpu / Primal 9.73025e-7 s 9.7125e-7 s 1.00
slicing / IDefOpt / tpu / Primal 0.0000010216 s 0.0000010199999999999998 s 1.00
slicing / JaXPipe / tpu / Forward 0.000001421 s 0.0000014013999999999998 s 1.01
slicing / Jax / tpu / Forward 0.000001476325 s 0.0000014771 s 1.00
slicing / HLOOpt / tpu / Forward 0.00000152175 s 0.000001513875 s 1.01
slicing / PartOpt / tpu / Forward 0.0000014932500000000005 s 0.000001488725 s 1.00
slicing / IPartOpt / tpu / Forward 0.0000015185499999999998 s 0.000001513225 s 1.00
slicing / DefOpt / tpu / Forward 0.0000014942 s 0.000001490975 s 1.00
slicing / IDefOpt / tpu / Forward 0.000001528575 s 0.000001517 s 1.01
slicing / JaXPipe / tpu / PreRev 0.00000257175 s 0.000002567925 s 1.00
slicing / JaXPipe / tpu / PostRev 0.00000252225 s 0.000002539775 s 0.99
slicing / JaXPipe / tpu / BothRev 0.000002589125 s 0.00000257955 s 1.00
slicing / Jax / tpu / BothRev 0.00000253295 s 0.0000025363 s 1.00
slicing / HLOOpt / tpu / PreRev 0.0000025906249999999994 s 0.00000258515 s 1.00
slicing / HLOOpt / tpu / PostRev 0.000002547325 s 0.0000025454500000000003 s 1.00
slicing / HLOOpt / tpu / BothRev 0.00000258375 s 0.000002587675 s 1.00
slicing / PartOpt / tpu / PreRev 0.000002536375 s 0.000002533125 s 1.00
slicing / PartOpt / tpu / PostRev 0.0000025854 s 0.00000257785 s 1.00
slicing / PartOpt / tpu / BothRev 0.000002544775000000001 s 0.00000254905 s 1.00
slicing / IPartOpt / tpu / PreRev 0.0000025765250000000003 s 0.000002587675 s 1.00
slicing / IPartOpt / tpu / PostRev 0.000002533075 s 0.000002545775 s 1.00
slicing / IPartOpt / tpu / BothRev 0.000002584 s 0.000002578575 s 1.00
slicing / DefOpt / tpu / PreRev 0.000002537525 s 0.000002541625 s 1.00
slicing / DefOpt / tpu / PostRev 0.000002590175 s 0.00000258455 s 1.00
slicing / DefOpt / tpu / BothRev 0.000002539125 s 0.0000025394250000000003 s 1.00
slicing / IDefOpt / tpu / PreRev 0.0000025740000000000003 s 0.00000257635 s 1.00
slicing / IDefOpt / tpu / PostRev 0.000002540225 s 0.0000025470500000000005 s 1.00
slicing / IDefOpt / tpu / BothRev 0.000002583975 s 0.000002589975 s 1.00
slicing / JaXPipe / cpu / Primal 0.000012551 s 0.0000065718999940145295 s 1.91
slicing / Jax / cpu / Primal 0.000012546 s 0.000006047020051482832 s 2.07
slicing / HLOOpt / cpu / Primal 0.000012641 s 0.000006182960005389759 s 2.04
slicing / PartOpt / cpu / Primal 0.000012438 s 0.000006284980026975973 s 1.98
slicing / IPartOpt / cpu / Primal 0.000012512 s 0.0000066760800291376655 s 1.87
slicing / DefOpt / cpu / Primal 0.000012751 s 0.000006011119994582259 s 2.12
slicing / IDefOpt / cpu / Primal 0.000012443 s 0.00000617611998677603 s 2.01
slicing / JaXPipe / cpu / Forward 0.000016907999999999998 s 0.00000992725998003152 s 1.70
slicing / Jax / cpu / Forward 0.000016804 s 0.000009392720012328936 s 1.79
slicing / HLOOpt / cpu / Forward 0.000016718 s 0.000010251100020468584 s 1.63
slicing / PartOpt / cpu / Forward 0.000016746000000000002 s 0.000009866499995041525 s 1.70
slicing / IPartOpt / cpu / Forward 0.000016778 s 0.000010019239998655394 s 1.67
slicing / DefOpt / cpu / Forward 0.000017010999999999997 s 0.000009773899992069344 s 1.74
slicing / IDefOpt / cpu / Forward 0.000016868000000000002 s 0.000009806559974094852 s 1.72
slicing / JaXPipe / cpu / PreRev 0.000017408 s 0.000010159940056837514 s 1.71
slicing / JaXPipe / cpu / PostRev 0.000017389999999999998 s 0.000010119239987034234 s 1.72
slicing / JaXPipe / cpu / BothRev 0.000017235 s 0.000010195520026172744 s 1.69
slicing / Jax / cpu / BothRev 0.000017462 s 0.000010721000016928885 s 1.63
slicing / HLOOpt / cpu / PreRev 0.000017643 s 0.000010868420013139255 s 1.62
slicing / HLOOpt / cpu / PostRev 0.000017526 s 0.000012423019970810856 s 1.41
slicing / HLOOpt / cpu / BothRev 0.000017279 s 0.000009850219994405052 s 1.75
slicing / PartOpt / cpu / PreRev 0.000017359 s 0.00000993288001154724 s 1.75
slicing / PartOpt / cpu / PostRev 0.000017754 s 0.000010278080017087632 s 1.73
slicing / PartOpt / cpu / BothRev 0.000017118 s 0.00001018304002172954 s 1.68
slicing / IPartOpt / cpu / PreRev 0.000017749999999999998 s 0.000010028739980043611 s 1.77
slicing / IPartOpt / cpu / PostRev 0.000018003 s 0.00001013477996821166 s 1.78
slicing / IPartOpt / cpu / BothRev 0.000017191000000000002 s 0.000009986599998228483 s 1.72
slicing / DefOpt / cpu / PreRev 0.00001747 s 0.000010073839976030283 s 1.73
slicing / DefOpt / cpu / PostRev 0.000017182 s 0.000010306819967809134 s 1.67
slicing / DefOpt / cpu / BothRev 0.000017506 s 0.000009928880026564 s 1.76
slicing / IDefOpt / cpu / PreRev 0.000017590999999999998 s 0.000010205140033576756 s 1.72
slicing / IDefOpt / cpu / PostRev 0.000017528 s 0.00001020074000734894 s 1.72
slicing / IDefOpt / cpu / BothRev 0.000017459 s 0.000010028259966929908 s 1.74
sum / JaXPipe / cpu / Primal 0.000008617719995527295 s 0.00000816919998214871 s 1.05
sum / Jax / cpu / Primal 0.00000831860005746421 s 0.000007457820038325735 s 1.12
sum / HLOOpt / cpu / Primal 0.00000850448005621729 s 0.00000763924002967542 s 1.11
sum / PartOpt / cpu / Primal 0.000008346859985977061 s 0.000007604439997521695 s 1.10
sum / IPartOpt / cpu / Primal 0.000008589820008637617 s 0.000008017280024432693 s 1.07
sum / DefOpt / cpu / Primal 0.000008474720007143332 s 0.000007534759988629957 s 1.12
sum / IDefOpt / cpu / Primal 0.00000825220000479021 s 0.000007843719968150254 s 1.05
sum / JaXPipe / cpu / Forward 0.000012614379957085477 s 0.0000113872000201809 s 1.11
sum / Jax / cpu / Forward 0.000012445560014384682 s 0.000011374180012353465 s 1.09
sum / HLOOpt / cpu / Forward 0.000012886859994978297 s 0.000011870960024680244 s 1.09
sum / PartOpt / cpu / Forward 0.000012763200002154918 s 0.000011473179984022864 s 1.11
sum / IPartOpt / cpu / Forward 0.000012605719994098764 s 0.000011139440020997428 s 1.13
sum / DefOpt / cpu / Forward 0.000012375119995340356 s 0.000011762800049837096 s 1.05
sum / IDefOpt / cpu / Forward 0.000012903699980597594 s 0.000011532799962878926 s 1.12
sum / JaXPipe / cpu / PreRev 0.000012399020024531637 s 0.000011403500038795756 s 1.09
sum / JaXPipe / cpu / PostRev 0.00001233142001183296 s 0.00001117342003453814 s 1.10
sum / JaXPipe / cpu / BothRev 0.000012024960024064057 s 0.000011444460023994906 s 1.05
sum / Jax / cpu / BothRev 0.000012290299973756192 s 0.00001101707997804624 s 1.12
sum / HLOOpt / cpu / PreRev 0.000012102540003979812 s 0.000011423160003687371 s 1.06
sum / HLOOpt / cpu / PostRev 0.000013930160012023407 s 0.000013029080000706018 s 1.07
sum / HLOOpt / cpu / BothRev 0.000012418359983712434 s 0.000010510639995118254 s 1.18
sum / PartOpt / cpu / PreRev 0.000012046579986417782 s 0.000011088860055679106 s 1.09
sum / PartOpt / cpu / PostRev 0.000012357460018392886 s 0.000011051279980165418 s 1.12
sum / PartOpt / cpu / BothRev 0.000012391380032568122 s 0.000010995619968525716 s 1.13
sum / IPartOpt / cpu / PreRev 0.00001221703999362944 s 0.000011126600020361364 s 1.10
sum / IPartOpt / cpu / PostRev 0.00001203181997880165 s 0.000011172479962624494 s 1.08
sum / IPartOpt / cpu / BothRev 0.000011863640020237654 s 0.000011434539965193835 s 1.04
sum / DefOpt / cpu / PreRev 0.000011986520057689633 s 0.00001090151998141664 s 1.10
sum / DefOpt / cpu / PostRev 0.00001207110002724221 s 0.000010847459998331031 s 1.11
sum / DefOpt / cpu / BothRev 0.000012361959998088425 s 0.00001036037999256223 s 1.19
sum / IDefOpt / cpu / PreRev 0.00001234378002664016 s 0.000010492980036360678 s 1.18
sum / IDefOpt / cpu / PostRev 0.000011972920028711086 s 0.000010710420028772205 s 1.12
sum / IDefOpt / cpu / BothRev 0.000012247439990460408 s 0.00001095347998671059 s 1.12
sum / JaXPipe / cuda / Primal 0.000002463 s 0.000002047 s 1.20
sum / Jax / cuda / Primal 0.000002464 s 0.000002047 s 1.20
sum / HLOOpt / cuda / Primal 0.000002463 s 0.000002047 s 1.20
sum / PartOpt / cuda / Primal 0.000002464 s 0.000002047 s 1.20
sum / IPartOpt / cuda / Primal 0.000002464 s 0.000002047 s 1.20
sum / DefOpt / cuda / Primal 0.000002464 s 0.000002047 s 1.20
sum / IDefOpt / cuda / Primal 0.000002464 s 0.000002047 s 1.20
sum / JaXPipe / cuda / Forward 0.000010592 s 0.000010336 s 1.02
sum / Jax / cuda / Forward 0.000010752 s 0.000010432 s 1.03
sum / HLOOpt / cuda / Forward 0.00001056 s 0.000009696 s 1.09
sum / PartOpt / cuda / Forward 0.0000104 s 0.00001024 s 1.02
sum / IPartOpt / cuda / Forward 0.000010687 s 0.0000096 s 1.11
sum / DefOpt / cuda / Forward 0.000010656 s 0.000010048 s 1.06
sum / IDefOpt / cuda / Forward 0.000010592 s 0.000010176 s 1.04
sum / JaXPipe / cuda / PreRev 0.00001072 s 0.000009791 s 1.09
sum / JaXPipe / cuda / PostRev 0.000010304 s 0.000009664 s 1.07
sum / JaXPipe / cuda / BothRev 0.000010208 s 0.000009823 s 1.04
sum / Jax / cuda / BothRev 0.000010335 s 0.000008928 s 1.16
sum / HLOOpt / cuda / PreRev 0.000010368 s 0.00000944 s 1.10
sum / HLOOpt / cuda / PostRev 0.000010496 s 0.000009728 s 1.08
sum / HLOOpt / cuda / BothRev 0.000010304 s 0.000009855 s 1.05
sum / PartOpt / cuda / PreRev 0.000011008 s 0.000010016 s 1.10
sum / PartOpt / cuda / PostRev 0.000010688 s 0.000009504 s 1.12
sum / PartOpt / cuda / BothRev 0.000010592 s 0.000009792 s 1.08
sum / IPartOpt / cuda / PreRev 0.000010656 s 0.00000976 s 1.09
sum / IPartOpt / cuda / PostRev 0.000010528 s 0.000010016 s 1.05
sum / IPartOpt / cuda / BothRev 0.000010976 s 0.000009824 s 1.12
sum / DefOpt / cuda / PreRev 0.000010688 s 0.000009824 s 1.09
sum / DefOpt / cuda / PostRev 0.0000104 s 0.000009568 s 1.09
sum / DefOpt / cuda / BothRev 0.00001088 s 0.000009472 s 1.15
sum / IDefOpt / cuda / PreRev 0.000010431 s 0.000009952 s 1.05
sum / IDefOpt / cuda / PostRev 0.000010368 s 0.000009728 s 1.07
sum / IDefOpt / cuda / BothRev 0.000011008 s 0.000009727 s 1.13
sum / JaXPipe / tpu / Primal 5.103e-7 s 5.102750000000001e-7 s 1.00
sum / Jax / tpu / Primal 5.470250000000001e-7 s 5.47075e-7 s 1.00
sum / HLOOpt / tpu / Primal 5.102750000000001e-7 s 5.106750000000001e-7 s 1.00
sum / PartOpt / tpu / Primal 5.473999999999999e-7 s 5.473e-7 s 1.00
sum / IPartOpt / tpu / Primal 5.108499999999999e-7 s 5.10375e-7 s 1.00
sum / DefOpt / tpu / Primal 5.4665e-7 s 5.467249999999999e-7 s 1.00
sum / IDefOpt / tpu / Primal 5.101e-7 s 5.104e-7 s 1.00
sum / JaXPipe / tpu / Forward 0.000001557675 s 0.0000015484 s 1.01
sum / Jax / tpu / Forward 0.0000015018000000000002 s 0.0000015049000000000002 s 1.00
sum / HLOOpt / tpu / Forward 0.00000153335 s 0.000001533725 s 1.00
sum / PartOpt / tpu / Forward 0.000001496625 s 0.0000014942 s 1.00
sum / IPartOpt / tpu / Forward 0.0000015390999999999998 s 0.0000015378249999999998 s 1.00
sum / DefOpt / tpu / Forward 0.000001496375 s 0.0000014951 s 1.00
sum / IDefOpt / tpu / Forward 0.0000015343 s 0.000001531225 s 1.00
sum / JaXPipe / tpu / PreRev 0.00000104545 s 0.0000010535 s 0.99
sum / JaXPipe / tpu / PostRev 0.0000010931250000000002 s 0.00000109215 s 1.00
sum / JaXPipe / tpu / BothRev 0.0000010528250000000002 s 0.000001053 s 1.00
sum / Jax / tpu / BothRev 0.000001084675 s 0.000001083075 s 1.00
sum / HLOOpt / tpu / PreRev 0.000001051875 s 0.000001052425 s 1.00
sum / HLOOpt / tpu / PostRev 0.000001089 s 0.00000109205 s 1.00
sum / HLOOpt / tpu / BothRev 0.000001052375 s 0.000001050175 s 1.00
sum / PartOpt / tpu / PreRev 0.000001087925 s 0.0000010928000000000002 s 1.00
sum / PartOpt / tpu / PostRev 0.0000010536 s 0.00000105045 s 1.00
sum / PartOpt / tpu / BothRev 0.0000010843 s 0.000001091575 s 0.99
sum / IPartOpt / tpu / PreRev 0.000001060125 s 0.00000105165 s 1.01
sum / IPartOpt / tpu / PostRev 0.0000011038 s 0.00000109235 s 1.01
sum / IPartOpt / tpu / BothRev 0.0000010480250000000002 s 0.000001055825 s 0.99
sum / DefOpt / tpu / PreRev 0.000001089975 s 0.0000010853 s 1.00
sum / DefOpt / tpu / PostRev 0.00000104725 s 0.00000104685 s 1.00
sum / DefOpt / tpu / BothRev 0.0000010896250000000002 s 0.000001083325 s 1.01
sum / IDefOpt / tpu / PreRev 0.0000010502250000000002 s 0.000001047075 s 1.00
sum / IDefOpt / tpu / PostRev 0.0000010848 s 0.000001092225 s 0.99
sum / IDefOpt / tpu / BothRev 0.00000105985 s 0.000001053925 s 1.01
sum / JaXPipe / cpu / Primal 0.000014935 s 0.00000816919998214871 s 1.83
sum / Jax / cpu / Primal 0.000014396 s 0.000007457820038325735 s 1.93
sum / HLOOpt / cpu / Primal 0.000014342 s 0.00000763924002967542 s 1.88
sum / PartOpt / cpu / Primal 0.000014296 s 0.000007604439997521695 s 1.88
sum / IPartOpt / cpu / Primal 0.000014648 s 0.000008017280024432693 s 1.83
sum / DefOpt / cpu / Primal 0.000014507 s 0.000007534759988629957 s 1.93
sum / IDefOpt / cpu / Primal 0.00001433 s 0.000007843719968150254 s 1.83
sum / JaXPipe / cpu / Forward 0.000020094 s 0.0000113872000201809 s 1.76
sum / Jax / cpu / Forward 0.000019542 s 0.000011374180012353465 s 1.72
sum / HLOOpt / cpu / Forward 0.00002836 s 0.000011870960024680244 s 2.39
sum / PartOpt / cpu / Forward 0.000020488 s 0.000011473179984022864 s 1.79
sum / IPartOpt / cpu / Forward 0.000019989 s 0.000011139440020997428 s 1.79
sum / DefOpt / cpu / Forward 0.000020216 s 0.000011762800049837096 s 1.72
sum / IDefOpt / cpu / Forward 0.000020444 s 0.000011532799962878926 s 1.77
sum / JaXPipe / cpu / PreRev 0.000019711 s 0.000011403500038795756 s 1.73
sum / JaXPipe / cpu / PostRev 0.000018731 s 0.00001117342003453814 s 1.68
sum / JaXPipe / cpu / BothRev 0.000018556 s 0.000011444460023994906 s 1.62
sum / Jax / cpu / BothRev 0.000019062 s 0.00001101707997804624 s 1.73
sum / HLOOpt / cpu / PreRev 0.000018692 s 0.000011423160003687371 s 1.64
sum / HLOOpt / cpu / PostRev 0.000018631 s 0.000013029080000706018 s 1.43
sum / HLOOpt / cpu / BothRev 0.000018392 s 0.000010510639995118254 s 1.75
sum / PartOpt / cpu / PreRev 0.000019293 s 0.000011088860055679106 s 1.74
sum / PartOpt / cpu / PostRev 0.000019012 s 0.000011051279980165418 s 1.72
sum / PartOpt / cpu / BothRev 0.00001887 s 0.000010995619968525716 s 1.72
sum / IPartOpt / cpu / PreRev 0.000018904 s 0.000011126600020361364 s 1.70
sum / IPartOpt / cpu / PostRev 0.000018773 s 0.000011172479962624494 s 1.68
sum / IPartOpt / cpu / BothRev 0.000018786 s 0.000011434539965193835 s 1.64
sum / DefOpt / cpu / PreRev 0.000018839 s 0.00001090151998141664 s 1.73
sum / DefOpt / cpu / PostRev 0.000018588 s 0.000010847459998331031 s 1.71
sum / DefOpt / cpu / BothRev 0.000018655 s 0.00001036037999256223 s 1.80
sum / IDefOpt / cpu / PreRev 0.000018509 s 0.000010492980036360678 s 1.76
sum / IDefOpt / cpu / PostRev 0.000019656 s 0.000010710420028772205 s 1.84
sum / IDefOpt / cpu / BothRev 0.000018677 s 0.00001095347998671059 s 1.71
value_and_grad / JaXPipe / cpu / Primal 0.000016022379977584933 s 0.000014257219991122838 s 1.12
value_and_grad / Jax / cpu / Primal 0.0000150683399533591 s 0.000014079739994485862 s 1.07
value_and_grad / HLOOpt / cpu / Primal 0.000015017479991001891 s 0.0000138554400200519 s 1.08
value_and_grad / PartOpt / cpu / Primal 0.00001524550002613978 s 0.000013732859997617196 s 1.11
value_and_grad / IPartOpt / cpu / Primal 0.000015583279991915334 s 0.000013589980017059134 s 1.15
value_and_grad / DefOpt / cpu / Primal 0.000015136359970711054 s 0.000013590080016001592 s 1.11
value_and_grad / IDefOpt / cpu / Primal 0.00001470868004616932 s 0.000013416999981927802 s 1.10
value_and_grad / JaXPipe / cuda / Primal 0.000034687 s 0.000033056 s 1.05
value_and_grad / Jax / cuda / Primal 0.000034592 s 0.000032705 s 1.06
value_and_grad / HLOOpt / cuda / Primal 0.000034688 s 0.000033119999999999995 s 1.05
value_and_grad / PartOpt / cuda / Primal 0.000034751 s 0.000033024 s 1.05
value_and_grad / IPartOpt / cuda / Primal 0.000038592 s 0.000032608 s 1.18
value_and_grad / DefOpt / cuda / Primal 0.000038304 s 0.000032288 s 1.19
value_and_grad / IDefOpt / cuda / Primal 0.000038623 s 0.000032864 s 1.18
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.000023509 s 0.000014257219991122838 s 1.65
value_and_grad / Jax / cpu / Primal 0.000022479 s 0.000014079739994485862 s 1.60
value_and_grad / HLOOpt / cpu / Primal 0.000023097 s 0.0000138554400200519 s 1.67
value_and_grad / PartOpt / cpu / Primal 0.000022676 s 0.000013732859997617196 s 1.65
value_and_grad / IPartOpt / cpu / Primal 0.000022903 s 0.000013589980017059134 s 1.69
value_and_grad / DefOpt / cpu / Primal 0.000023155 s 0.000013590080016001592 s 1.70
value_and_grad / IDefOpt / cpu / Primal 0.000023462 s 0.000013416999981927802 s 1.75
jaxmd20 / JaXPipe / cuda / Primal 0.001458259 s 0.001457397 s 1.00
jaxmd20 / Jax / cuda / Primal 0.001547825 s 0.001451123 s 1.07
jaxmd20 / HLOOpt / cuda / Primal 0.001348116 s 0.001325206 s 1.02
jaxmd20 / PartOpt / cuda / Primal 0.0013803069999999 s 0.00132207 s 1.04
jaxmd20 / IPartOpt / cuda / Primal 0.001343412 s 0.001354484 s 0.99
jaxmd20 / DefOpt / cuda / Primal 0.000939959 s 0.000915832 s 1.03
jaxmd20 / IDefOpt / cuda / Primal 0.000963319 s 0.000946264 s 1.02
jaxmd20 / JaXPipe / cuda / Forward 0.001629169 s 0.001565299 s 1.04
jaxmd20 / Jax / cuda / Forward 0.0019210689999999 s 0.00177445 s 1.08
jaxmd20 / HLOOpt / cuda / Forward 0.001711919 s 0.0016276029999999 s 1.05
jaxmd20 / PartOpt / cuda / Forward 0.0017235039999999 s 0.001644435 s 1.05
jaxmd20 / IPartOpt / cuda / Forward 0.00170632 s 0.001633073 s 1.04
jaxmd20 / DefOpt / cuda / Forward 0.001712944 s 0.001641395 s 1.04
jaxmd20 / IDefOpt / cuda / Forward 0.001710736 s 0.001626067 s 1.05
jaxmd20 / JaXPipe / cuda / PreRev 0.002757317 s 0.002650986 s 1.04
jaxmd20 / JaXPipe / cuda / PostRev 0.005418795 s 0.005307668 s 1.02
jaxmd20 / JaXPipe / cuda / BothRev 0.0027422469999999 s 0.002674538 s 1.03
jaxmd20 / Jax / cuda / BothRev 0.005448206 s 0.005279286 s 1.03
jaxmd20 / HLOOpt / cuda / PreRev 0.0028685779999999 s 0.002728745 s 1.05
jaxmd20 / HLOOpt / cuda / PostRev 0.005482765 s 0.00529202 s 1.04
jaxmd20 / HLOOpt / cuda / BothRev 0.0028035869999999 s 0.002721225 s 1.03
jaxmd20 / PartOpt / cuda / PreRev 0.002888548 s 0.002842761 s 1.02
jaxmd20 / PartOpt / cuda / PostRev 0.005585579 s 0.005400819 s 1.03
jaxmd20 / PartOpt / cuda / BothRev 0.002839462 s 0.002791464 s 1.02
jaxmd20 / IPartOpt / cuda / PreRev 0.0029305 s 0.002799721 s 1.05
jaxmd20 / IPartOpt / cuda / PostRev 0.005556685 s 0.005345107 s 1.04
jaxmd20 / IPartOpt / cuda / BothRev 0.002898244 s 0.0027779289999999 s 1.04
jaxmd20 / DefOpt / cuda / PreRev 0.002917028 s 0.0028215439999999 s 1.03
jaxmd20 / DefOpt / cuda / PostRev 0.002827398 s 0.002743369 s 1.03
jaxmd20 / DefOpt / cuda / BothRev 0.002825992 s 0.002758602 s 1.02
jaxmd20 / IDefOpt / cuda / PreRev 0.002913796 s 0.002804522 s 1.04
jaxmd20 / IDefOpt / cuda / PostRev 0.002379273 s 0.00230478 s 1.03
jaxmd20 / IDefOpt / cuda / BothRev 0.002825861 s 0.002734473 s 1.03
jaxmd20 / JaXPipe / tpu / Primal 0.009263875 s 0.009273954375 s 1.00
jaxmd20 / Jax / tpu / Primal 0.009277950625 s 0.00926473125 s 1.00
jaxmd20 / HLOOpt / tpu / Primal 0.009164131875 s 0.009170430625 s 1.00
jaxmd20 / PartOpt / tpu / Primal 0.009201443125 s 0.00919640125 s 1.00
jaxmd20 / IPartOpt / tpu / Primal 0.009202941875 s 0.0092012087499999 s 1.00
jaxmd20 / DefOpt / tpu / Primal 0.008805996875 s 0.0087924475 s 1.00
jaxmd20 / IDefOpt / tpu / Primal 0.0086989593749999 s 0.00869956375 s 1.00
jaxmd20 / JaXPipe / tpu / Forward 0.017419074375 s 0.0174177849999999 s 1.00
jaxmd20 / Jax / tpu / Forward 0.018725729375 s 0.018727536875 s 1.00
jaxmd20 / HLOOpt / tpu / Forward 0.017395253125 s 0.0173932675 s 1.00
jaxmd20 / PartOpt / tpu / Forward 0.0174224825 s 0.017409841875 s 1.00
jaxmd20 / IPartOpt / tpu / Forward 0.017417525 s 0.01741184875 s 1.00
jaxmd20 / DefOpt / tpu / Forward 0.0174133475 s 0.017417063125 s 1.00
jaxmd20 / IDefOpt / tpu / Forward 0.0174168718749999 s 0.017414838125 s 1.00
jaxmd20 / JaXPipe / tpu / PreRev 0.02546077 s 0.025457535 s 1.00
jaxmd20 / JaXPipe / tpu / PostRev 0.0218534225 s 0.0218942125 s 1.00
jaxmd20 / JaXPipe / tpu / BothRev 0.025467985625 s 0.02547417375 s 1.00
jaxmd20 / Jax / tpu / BothRev 0.021855240625 s 0.021891351875 s 1.00
jaxmd20 / HLOOpt / tpu / PreRev 0.025586429375 s 0.0255871593749999 s 1.00
jaxmd20 / HLOOpt / tpu / PostRev 0.0208013024999999 s 0.02083006625 s 1.00
jaxmd20 / HLOOpt / tpu / BothRev 0.02568960875 s 0.025685535 s 1.00
jaxmd20 / PartOpt / tpu / PreRev 0.025441500625 s 0.02550769125 s 1.00
jaxmd20 / PartOpt / tpu / PostRev 0.021513925625 s 0.02150850125 s 1.00
jaxmd20 / PartOpt / tpu / BothRev 0.025532771875 s 0.025594076875 s 1.00
jaxmd20 / IPartOpt / tpu / PreRev 0.025472556875 s 0.025477834375 s 1.00
jaxmd20 / IPartOpt / tpu / PostRev 0.021506566875 s 0.021536304375 s 1.00
jaxmd20 / IPartOpt / tpu / BothRev 0.025560885625 s 0.025555453125 s 1.00
jaxmd20 / DefOpt / tpu / PreRev 0.025441096875 s 0.025506061875 s 1.00
jaxmd20 / DefOpt / tpu / PostRev 0.01881985375 s 0.01880455625 s 1.00
jaxmd20 / DefOpt / tpu / BothRev 0.02553337875 s 0.02559681875 s 1.00
jaxmd20 / IDefOpt / tpu / PreRev 0.025473979375 s 0.025477278125 s 1.00
jaxmd20 / IDefOpt / tpu / PostRev 0.018296206875 s 0.01834256875 s 1.00
jaxmd20 / IDefOpt / tpu / BothRev 0.025564598125 s 0.025552339375 s 1.00
jaxmd40 / JaXPipe / cpu / Primal 0.06017587 s 0.080197468 s 0.75
jaxmd40 / Jax / cpu / Primal 0.067363171 s 0.083655234 s 0.81
jaxmd40 / HLOOpt / cpu / Primal 0.090845755 s 0.096154745 s 0.94
jaxmd40 / PartOpt / cpu / Primal 0.071163119 s 0.083441797 s 0.85
jaxmd40 / IPartOpt / cpu / Primal 0.071862487 s 0.07348125 s 0.98
jaxmd40 / DefOpt / cpu / Primal 0.087342267 s 0.1080883209999999 s 0.81
jaxmd40 / IDefOpt / cpu / Primal 0.088585051 s 0.102355934 s 0.87
jaxmd40 / JaXPipe / cpu / Forward 0.154162945 s 0.194925377 s 0.79
jaxmd40 / Jax / cpu / Forward 0.078697899 s 0.102705963 s 0.77
jaxmd40 / HLOOpt / cpu / Forward 0.153707904 s 0.1943103319999999 s 0.79
jaxmd40 / PartOpt / cpu / Forward 0.155135194 s 0.192890015 s 0.80
jaxmd40 / IPartOpt / cpu / Forward 0.158978913 s 0.1889543649999999 s 0.84
jaxmd40 / DefOpt / cpu / Forward 0.158821388 s 0.187598506 s 0.85
jaxmd40 / IDefOpt / cpu / Forward 0.158957989 s 0.1914111 s 0.83
jaxmd40 / JaXPipe / cpu / PreRev 0.23184539 s 0.2570904029999999 s 0.90
jaxmd40 / JaXPipe / cpu / PostRev 0.138482028 s 0.164598411 s 0.84
jaxmd40 / JaXPipe / cpu / BothRev 0.2245381769999999 s 0.243595934 s 0.92
jaxmd40 / Jax / cpu / BothRev 0.123784988 s 0.167637033 s 0.74
jaxmd40 / HLOOpt / cpu / PreRev 0.237255857 s 0.247185092 s 0.96
jaxmd40 / HLOOpt / cpu / PostRev 0.1653489 s 0.214137107 s 0.77
jaxmd40 / HLOOpt / cpu / BothRev 0.253492825 s 0.28978097 s 0.87
jaxmd40 / PartOpt / cpu / PreRev 0.239485215 s 0.245707196 s 0.97
jaxmd40 / PartOpt / cpu / PostRev 0.1431804119999999 s 0.148089955 s 0.97
jaxmd40 / PartOpt / cpu / BothRev 0.266242319 s 0.289875612 s 0.92
jaxmd40 / IPartOpt / cpu / PreRev 0.224458415 s 0.250292784 s 0.90
jaxmd40 / IPartOpt / cpu / PostRev 0.127183482 s 0.155665586 s 0.82
jaxmd40 / IPartOpt / cpu / BothRev 0.2411383889999999 s 0.280760357 s 0.86
jaxmd40 / DefOpt / cpu / PreRev 0.217700555 s 0.246668072 s 0.88
jaxmd40 / DefOpt / cpu / PostRev 0.175047671 s 0.208282383 s 0.84
jaxmd40 / DefOpt / cpu / BothRev 0.246552671 s 0.286270934 s 0.86
jaxmd40 / IDefOpt / cpu / PreRev 0.208378323 s 0.243850494 s 0.85
jaxmd40 / IDefOpt / cpu / PostRev 0.170986614 s 0.2124084 s 0.80
jaxmd40 / IDefOpt / cpu / BothRev 0.246931089 s 0.2716668399999999 s 0.91
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / cuda / Primal 1.706815601 s 1.7028338120000002 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / cuda / Primal 1.709699376 s 1.704407339 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / cuda / Primal 1.7230681790000002 s 1.715157461 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / cuda / Primal 1.7013737 s 1.695699103 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / cuda / Primal 1.699425423 s 1.6934279449999998 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / cuda / Primal 1.6703023069999998 s 1.664433072 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / cuda / Primal 1.925637318 s 1.921661585 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / tpu / Primal 3.038602243125 s 3.038416416875 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / tpu / Primal 3.03918639125 s 3.038972848125 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / tpu / Primal 3.121468010625 s 3.12138719375 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / tpu / Primal 3.0599137137500003 s 3.059688449375 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / tpu / Primal 3.0602407775 s 3.05992763375 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / tpu / Primal 2.102394579375 s 2.102344878125 s 1.00
v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / tpu / Primal 2.944494238125 s 2.944463474375 s 1.00
v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / JaXPipe / cpu / Primal 6.124516997 s 6.8007833060000005 s 0.90
v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / Jax / cpu / Primal 5.970445896 s 6.818972752 s 0.88
v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / HLOOpt / cpu / Primal 5.985459457 s 6.647598244 s 0.90
v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / PartOpt / cpu / Primal 6.241103656 s 6.899143466000001 s 0.90
v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IPartOpt / cpu / Primal 6.136173479 s 6.961332216 s 0.88
v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / DefOpt / cpu / Primal 2.3618552420000003 s 2.786819339 s 0.85
v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IDefOpt / cpu / Primal 6.526824231 s 7.423087524 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