Skip to content

feat: expand mul scatter simply to support setindex ones #1888

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
Merged
avik-pal merged 2 commits into from
Jan 6, 2026
Merged

feat: expand mul scatter simply to support setindex ones #1888

avik-pal merged 2 commits into from
Jan 6, 2026

Conversation

@avik-pal
Copy link
Collaborator

@avik-pal avik-pal commented Jan 5, 2026

No description provided.

@avik-pal avik-pal requested a review from wsmoses January 5, 2026 22:55
@avik-pal
Copy link
Collaborator Author

avik-pal commented Jan 6, 2026

@wsmoses mac extendrotate seems broken

@avik-pal avik-pal mentioned this pull request Jan 6, 2026
Open
@avik-pal avik-pal merged commit a83eae4 into main Jan 6, 2026
24 of 25 checks passed
@avik-pal avik-pal deleted the ap/more_scatter_simplify branch January 6, 2026 03:12
github-actions[bot]
github-actions bot reviewed Jan 6, 2026
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

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

EnzymeJAX Benchmarks

Details
Benchmark suite Current: ea59b14 Previous: 972c249 Ratio
actmtch / JaXPipe / cpu / Primal 0.000006859219993202714 s 0.000006562319977092557 s 1.05
actmtch / Jax / cpu / Primal 0.000006495100005849964 s 0.0000062876599804440046 s 1.03
actmtch / HLOOpt / cpu / Primal 0.000007429160004903678 s 0.000007609099993715063 s 0.98
actmtch / PartOpt / cpu / Primal 0.000006361680011650605 s 0.000006221880139491986 s 1.02
actmtch / IPartOpt / cpu / Primal 0.000006740979997630347 s 0.000006396459939423949 s 1.05
actmtch / DefOpt / cpu / Primal 0.000006959980005376565 s 0.000007537180008512223 s 0.92
actmtch / IDefOpt / cpu / Primal 0.00000684888000250794 s 0.000006963640062167542 s 0.98
actmtch / JaXPipe / cpu / Forward 0.000010989500008236064 s 0.00001072302004104131 s 1.02
actmtch / Jax / cpu / Forward 0.000010224340010154264 s 0.000010128840040124487 s 1.01
actmtch / HLOOpt / cpu / Forward 0.00001141735999908633 s 0.000010821760133694624 s 1.06
actmtch / PartOpt / cpu / Forward 0.000010710779993132746 s 0.00001043737996951677 s 1.03
actmtch / IPartOpt / cpu / Forward 0.000011185579996890736 s 0.000010789780044433427 s 1.04
actmtch / DefOpt / cpu / Forward 0.000010735419996308338 s 0.0000100856801327609 s 1.06
actmtch / IDefOpt / cpu / Forward 0.000011227639995468052 s 0.00001044433998686145 s 1.07
actmtch / JaXPipe / cpu / PreRev 0.000010866119998809154 s 0.00001114240005335887 s 0.98
actmtch / JaXPipe / cpu / PostRev 0.00001084380000065721 s 0.000009797399943636265 s 1.11
actmtch / JaXPipe / cpu / BothRev 0.000011416160009503074 s 0.00001141130011092173 s 1.00
actmtch / Jax / cpu / BothRev 0.000009570439995059132 s 0.000009590319968992844 s 1.00
actmtch / HLOOpt / cpu / PreRev 0.00001137167999331723 s 0.00001094651990570128 s 1.04
actmtch / HLOOpt / cpu / PostRev 0.00001385724000101618 s 0.000012295760061533656 s 1.13
actmtch / HLOOpt / cpu / BothRev 0.000011121520001324824 s 0.000010983460069837748 s 1.01
actmtch / PartOpt / cpu / PreRev 0.000011051760004647804 s 0.000010270799921272557 s 1.08
actmtch / PartOpt / cpu / PostRev 0.000010784779988171069 s 0.000009451600017200687 s 1.14
actmtch / PartOpt / cpu / BothRev 0.000011583279988371942 s 0.000011059059979743323 s 1.05
actmtch / IPartOpt / cpu / PreRev 0.000010450619995481248 s 0.000010759480028355029 s 0.97
actmtch / IPartOpt / cpu / PostRev 0.000009947260007265868 s 0.000009635619935579598 s 1.03
actmtch / IPartOpt / cpu / BothRev 0.000011062040009619522 s 0.000010833199994522146 s 1.02
actmtch / DefOpt / cpu / PreRev 0.000010863479988074686 s 0.000010563840023678494 s 1.03
actmtch / DefOpt / cpu / PostRev 0.00001130106000118758 s 0.000010677940026653233 s 1.06
actmtch / DefOpt / cpu / BothRev 0.000011159060004501952 s 0.000011130339989904314 s 1.00
actmtch / IDefOpt / cpu / PreRev 0.000010534359998928269 s 0.000010484219928912352 s 1.00
actmtch / IDefOpt / cpu / PostRev 0.000011421360009080672 s 0.00001111930003389716 s 1.03
actmtch / IDefOpt / cpu / BothRev 0.00001143517999935284 s 0.000010797219929372658 s 1.06
actmtch / JaXPipe / cuda / Primal 0.000002015 s 0.0000024 s 0.84
actmtch / Jax / cuda / Primal 0.000002015 s 0.000002399 s 0.84
actmtch / HLOOpt / cuda / Primal 0.000002016 s 0.0000024 s 0.84
actmtch / PartOpt / cuda / Primal 0.000002016 s 0.0000024 s 0.84
actmtch / IPartOpt / cuda / Primal 0.000002016 s 0.0000024 s 0.84
actmtch / DefOpt / cuda / Primal 0.000002015 s 0.0000024 s 0.84
actmtch / IDefOpt / cuda / Primal 0.000002016 s 0.000002399 s 0.84
actmtch / JaXPipe / cuda / Forward 0.000009824 s 0.00001024 s 0.96
actmtch / Jax / cuda / Forward 0.000010465 s 0.000010432 s 1.00
actmtch / HLOOpt / cuda / Forward 0.000010048 s 0.000010176 s 0.99
actmtch / PartOpt / cuda / Forward 0.000010048 s 0.000010592 s 0.95
actmtch / IPartOpt / cuda / Forward 0.000010304 s 0.000010303 s 1.00
actmtch / DefOpt / cuda / Forward 0.000009984 s 0.000010496 s 0.95
actmtch / IDefOpt / cuda / Forward 0.000010144 s 0.000010369 s 0.98
actmtch / JaXPipe / cuda / PreRev 0.000010176 s 0.000010784 s 0.94
actmtch / JaXPipe / cuda / PostRev 0.00001008 s 0.000010624 s 0.95
actmtch / JaXPipe / cuda / BothRev 0.000010304 s 0.000010368 s 0.99
actmtch / Jax / cuda / BothRev 0.000010272 s 0.000010624 s 0.97
actmtch / HLOOpt / cuda / PreRev 0.0000104 s 0.0000104 s 1
actmtch / HLOOpt / cuda / PostRev 0.00001024 s 0.0000104 s 0.98
actmtch / HLOOpt / cuda / BothRev 0.000010656 s 0.000010336 s 1.03
actmtch / PartOpt / cuda / PreRev 0.000010272 s 0.000010463 s 0.98
actmtch / PartOpt / cuda / PostRev 0.000010144 s 0.00001056 s 0.96
actmtch / PartOpt / cuda / BothRev 0.000010336 s 0.000010592 s 0.98
actmtch / IPartOpt / cuda / PreRev 0.000010368 s 0.0000104 s 1.00
actmtch / IPartOpt / cuda / PostRev 0.000010399 s 0.000010369 s 1.00
actmtch / IPartOpt / cuda / BothRev 0.000010272 s 0.000011232 s 0.91
actmtch / DefOpt / cuda / PreRev 0.000010176 s 0.000011616 s 0.88
actmtch / DefOpt / cuda / PostRev 0.000010272 s 0.00001136 s 0.90
actmtch / DefOpt / cuda / BothRev 0.00001024 s 0.00001072 s 0.96
actmtch / IDefOpt / cuda / PreRev 0.000010272 s 0.000010432 s 0.98
actmtch / IDefOpt / cuda / PostRev 0.000010016 s 0.00001072 s 0.93
actmtch / IDefOpt / cuda / BothRev 0.000010208 s 0.000010368 s 0.98
actmtch / JaXPipe / tpu / Primal 5.82675e-7 s 5.633749999999999e-7 s 1.03
actmtch / Jax / tpu / Primal 5.6365e-7 s 5.973e-7 s 0.94
actmtch / HLOOpt / tpu / Primal 0.00000216795 s 0.000002102425 s 1.03
actmtch / PartOpt / tpu / Primal 5.63675e-7 s 5.96875e-7 s 0.94
actmtch / IPartOpt / tpu / Primal 5.755749999999999e-7 s 5.5285e-7 s 1.04
actmtch / DefOpt / tpu / Primal 0.00000205265 s 0.00000215745 s 0.95
actmtch / IDefOpt / tpu / Primal 0.0000021682000000000003 s 0.000002101375 s 1.03
actmtch / JaXPipe / tpu / Forward 0.000003871675 s 0.000003832575 s 1.01
actmtch / Jax / tpu / Forward 0.000001244425 s 0.000001213375 s 1.03
actmtch / HLOOpt / tpu / Forward 0.00000365765 s 0.0000039319250000000005 s 0.93
actmtch / PartOpt / tpu / Forward 0.000003908975 s 0.00000391155 s 1.00
actmtch / IPartOpt / tpu / Forward 0.000003665525 s 0.000003945575 s 0.93
actmtch / DefOpt / tpu / Forward 0.000003882825 s 0.000003906125 s 0.99
actmtch / IDefOpt / tpu / Forward 0.000003680875 s 0.0000039276 s 0.94
actmtch / JaXPipe / tpu / PreRev 0.00000375075 s 0.000003482075 s 1.08
actmtch / JaXPipe / tpu / PostRev 0.0000016194 s 0.0000016388 s 0.99
actmtch / JaXPipe / tpu / BothRev 0.000003757475 s 0.00000348365 s 1.08
actmtch / Jax / tpu / BothRev 0.0000016265 s 0.0000016362 s 0.99
actmtch / HLOOpt / tpu / PreRev 0.000003747175 s 0.0000034763 s 1.08
actmtch / HLOOpt / tpu / PostRev 0.0000034604000000000004 s 0.0000034209750000000003 s 1.01
actmtch / HLOOpt / tpu / BothRev 0.0000037489 s 0.000003486025 s 1.08
actmtch / PartOpt / tpu / PreRev 0.00000344645 s 0.000003394725 s 1.02
actmtch / PartOpt / tpu / PostRev 0.0000016681 s 0.0000015865000000000002 s 1.05
actmtch / PartOpt / tpu / BothRev 0.000003453775 s 0.000003427625 s 1.01
actmtch / IPartOpt / tpu / PreRev 0.000003747 s 0.000003497 s 1.07
actmtch / IPartOpt / tpu / PostRev 0.0000016234750000000002 s 0.0000016344 s 0.99
actmtch / IPartOpt / tpu / BothRev 0.000003742425 s 0.0000034946 s 1.07
actmtch / DefOpt / tpu / PreRev 0.00000344925 s 0.000003425075 s 1.01
actmtch / DefOpt / tpu / PostRev 0.0000036825 s 0.000003411675 s 1.08
actmtch / DefOpt / tpu / BothRev 0.000003451725 s 0.000003417575 s 1.01
actmtch / IDefOpt / tpu / PreRev 0.0000037407 s 0.00000347555 s 1.08
actmtch / IDefOpt / tpu / PostRev 0.0000034393249999999995 s 0.000003417775 s 1.01
actmtch / IDefOpt / tpu / BothRev 0.00000375055 s 0.0000034766249999999995 s 1.08
actmtch / JaXPipe / cpu / Primal 0.000016413 s 0.000006562319977092557 s 2.50
actmtch / Jax / cpu / Primal 0.000021564 s 0.0000062876599804440046 s 3.43
actmtch / HLOOpt / cpu / Primal 0.000017657 s 0.000007609099993715063 s 2.32
actmtch / PartOpt / cpu / Primal 0.000016417999999999998 s 0.000006221880139491986 s 2.64
actmtch / IPartOpt / cpu / Primal 0.000016292999999999998 s 0.000006396459939423949 s 2.55
actmtch / DefOpt / cpu / Primal 0.000017612 s 0.000007537180008512223 s 2.34
actmtch / IDefOpt / cpu / Primal 0.000017357 s 0.000006963640062167542 s 2.49
actmtch / JaXPipe / cpu / Forward 0.000024529 s 0.00001072302004104131 s 2.29
actmtch / Jax / cpu / Forward 0.000022219 s 0.000010128840040124487 s 2.19
actmtch / HLOOpt / cpu / Forward 0.000023648 s 0.000010821760133694624 s 2.19
actmtch / PartOpt / cpu / Forward 0.000023845000000000003 s 0.00001043737996951677 s 2.28
actmtch / IPartOpt / cpu / Forward 0.000023835 s 0.000010789780044433427 s 2.21
actmtch / DefOpt / cpu / Forward 0.000023695 s 0.0000100856801327609 s 2.35
actmtch / IDefOpt / cpu / Forward 0.000024015000000000003 s 0.00001044433998686145 s 2.30
actmtch / JaXPipe / cpu / PreRev 0.000024415 s 0.00001114240005335887 s 2.19
actmtch / JaXPipe / cpu / PostRev 0.000022045 s 0.000009797399943636265 s 2.25
actmtch / JaXPipe / cpu / BothRev 0.000024402 s 0.00001141130011092173 s 2.14
actmtch / Jax / cpu / BothRev 0.000022034000000000003 s 0.000009590319968992844 s 2.30
actmtch / HLOOpt / cpu / PreRev 0.00002419 s 0.00001094651990570128 s 2.21
actmtch / HLOOpt / cpu / PostRev 0.000023882 s 0.000012295760061533656 s 1.94
actmtch / HLOOpt / cpu / BothRev 0.000024318 s 0.000010983460069837748 s 2.21
actmtch / PartOpt / cpu / PreRev 0.000023818 s 0.000010270799921272557 s 2.32
actmtch / PartOpt / cpu / PostRev 0.000022403 s 0.000009451600017200687 s 2.37
actmtch / PartOpt / cpu / BothRev 0.00002429 s 0.000011059059979743323 s 2.20
actmtch / IPartOpt / cpu / PreRev 0.000023751 s 0.000010759480028355029 s 2.21
actmtch / IPartOpt / cpu / PostRev 0.000021962 s 0.000009635619935579598 s 2.28
actmtch / IPartOpt / cpu / BothRev 0.000024075 s 0.000010833199994522146 s 2.22
actmtch / DefOpt / cpu / PreRev 0.000023875 s 0.000010563840023678494 s 2.26
actmtch / DefOpt / cpu / PostRev 0.000024587 s 0.000010677940026653233 s 2.30
actmtch / DefOpt / cpu / BothRev 0.000023846 s 0.000011130339989904314 s 2.14
actmtch / IDefOpt / cpu / PreRev 0.000024074 s 0.000010484219928912352 s 2.30
actmtch / IDefOpt / cpu / PostRev 0.000023998 s 0.00001111930003389716 s 2.16
actmtch / IDefOpt / cpu / BothRev 0.000024738 s 0.000010797219929372658 s 2.29
add_one / JaXPipe / cpu / Primal 0.0000067783000054078005 s 0.000006647580121352803 s 1.02
add_one / Jax / cpu / Primal 0.000006589159984287107 s 0.000007019260065135313 s 0.94
add_one / HLOOpt / cpu / Primal 0.000006782279997423757 s 0.000006855379924672889 s 0.99
add_one / PartOpt / cpu / Primal 0.0000066239399848200266 s 0.000006245519980438985 s 1.06
add_one / IPartOpt / cpu / Primal 0.000007073600006606284 s 0.000006942599975445774 s 1.02
add_one / DefOpt / cpu / Primal 0.000006496879991573223 s 0.0000069157200050540265 s 0.94
add_one / IDefOpt / cpu / Primal 0.000006371919998855446 s 0.00000644126001134282 s 0.99
add_one / JaXPipe / cpu / Forward 0.000010277220003445108 s 0.000010058619955088945 s 1.02
add_one / Jax / cpu / Forward 0.000010390500003722993 s 0.000009927860028255965 s 1.05
add_one / HLOOpt / cpu / Forward 0.00001029220001100839 s 0.000010169480028707766 s 1.01
add_one / PartOpt / cpu / Forward 0.000010512819992527513 s 0.000010235980043944435 s 1.03
add_one / IPartOpt / cpu / Forward 0.000010377800003880111 s 0.00000981865994617692 s 1.06
add_one / DefOpt / cpu / Forward 0.00001015829999687412 s 0.000010122280018549644 s 1.00
add_one / IDefOpt / cpu / Forward 0.000010026439997545824 s 0.000009593060021870769 s 1.05
add_one / JaXPipe / cpu / PreRev 0.000011822139995274484 s 0.000011487320007290693 s 1.03
add_one / JaXPipe / cpu / PostRev 0.000011467800002264994 s 0.000011399719842302149 s 1.01
add_one / JaXPipe / cpu / BothRev 0.000011673679996420103 s 0.000011612500002229353 s 1.01
add_one / Jax / cpu / BothRev 0.00001170005999483692 s 0.000011088899955211672 s 1.06
add_one / HLOOpt / cpu / PreRev 0.000011623739997048688 s 0.000011683019965857968 s 0.99
add_one / HLOOpt / cpu / PostRev 0.000013904540007843024 s 0.000014739599992026342 s 0.94
add_one / HLOOpt / cpu / BothRev 0.00001099073999284883 s 0.000011134360011055832 s 0.99
add_one / PartOpt / cpu / PreRev 0.00001161570000476786 s 0.0000115139000081399 s 1.01
add_one / PartOpt / cpu / PostRev 0.000011506360001476424 s 0.000011024920004274463 s 1.04
add_one / PartOpt / cpu / BothRev 0.00001234815999168859 s 0.00001160763993539149 s 1.06
add_one / IPartOpt / cpu / PreRev 0.000011284499996691012 s 0.000011384020108380355 s 0.99
add_one / IPartOpt / cpu / PostRev 0.000011400479993426417 s 0.000011333639959048014 s 1.01
add_one / IPartOpt / cpu / BothRev 0.000011690819992509204 s 0.000010866900011023972 s 1.08
add_one / DefOpt / cpu / PreRev 0.000011831100002837048 s 0.00001100651994420332 s 1.07
add_one / DefOpt / cpu / PostRev 0.000011357640003097912 s 0.000010951919994113267 s 1.04
add_one / DefOpt / cpu / BothRev 0.00001120451999895522 s 0.000010872719994949876 s 1.03
add_one / IDefOpt / cpu / PreRev 0.000011536099996192209 s 0.000011335459948895732 s 1.02
add_one / IDefOpt / cpu / PostRev 0.00001187340001024495 s 0.000010873259907384635 s 1.09
add_one / IDefOpt / cpu / BothRev 0.000011828839992631402 s 0.000010865979947993764 s 1.09
add_one / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.000002304 s 0.83
add_one / Jax / cuda / Primal 0.0000019200000000000003 s 0.000002335 s 0.82
add_one / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.000002303 s 0.83
add_one / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.000002335 s 0.82
add_one / IPartOpt / cuda / Primal 0.000001919 s 0.000002335 s 0.82
add_one / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.000002335 s 0.82
add_one / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.000002335 s 0.82
add_one / JaXPipe / cuda / Forward 0.000010272 s 0.000010752 s 0.96
add_one / Jax / cuda / Forward 0.000010272 s 0.000010432 s 0.98
add_one / HLOOpt / cuda / Forward 0.000010048 s 0.000010592 s 0.95
add_one / PartOpt / cuda / Forward 0.00000992 s 0.000013215 s 0.75
add_one / IPartOpt / cuda / Forward 0.000010304 s 0.000010304 s 1
add_one / DefOpt / cuda / Forward 0.000010335 s 0.000010368 s 1.00
add_one / IDefOpt / cuda / Forward 0.000010304 s 0.000010175 s 1.01
add_one / JaXPipe / cuda / PreRev 0.000025376 s 0.000025792 s 0.98
add_one / JaXPipe / cuda / PostRev 0.000025472000000000003 s 0.000025536 s 1.00
add_one / JaXPipe / cuda / BothRev 0.000025824 s 0.000025696 s 1.00
add_one / Jax / cuda / BothRev 0.000025664 s 0.000025792 s 1.00
add_one / HLOOpt / cuda / PreRev 0.0000256 s 0.000025568 s 1.00
add_one / HLOOpt / cuda / PostRev 0.000025184 s 0.000024256 s 1.04
add_one / HLOOpt / cuda / BothRev 0.000026304 s 0.000025568 s 1.03
add_one / PartOpt / cuda / PreRev 0.00002592 s 0.000025376 s 1.02
add_one / PartOpt / cuda / PostRev 0.0000256 s 0.00002576 s 0.99
add_one / PartOpt / cuda / BothRev 0.000025344 s 0.000025632 s 0.99
add_one / IPartOpt / cuda / PreRev 0.00002608 s 0.000026272 s 0.99
add_one / IPartOpt / cuda / PostRev 0.000026176 s 0.000025409 s 1.03
add_one / IPartOpt / cuda / BothRev 0.00002544 s 0.000025919 s 0.98
add_one / DefOpt / cuda / PreRev 0.000025952 s 0.00002608 s 1.00
add_one / DefOpt / cuda / PostRev 0.000025504 s 0.000025695 s 0.99
add_one / DefOpt / cuda / BothRev 0.00002512 s 0.000026016 s 0.97
add_one / IDefOpt / cuda / PreRev 0.00002624 s 0.0000248 s 1.06
add_one / IDefOpt / cuda / PostRev 0.000025695 s 0.000025888 s 0.99
add_one / IDefOpt / cuda / BothRev 0.000026592 s 0.000025504 s 1.04
add_one / JaXPipe / tpu / Primal 0.0000014493500000000002 s 0.00000142345 s 1.02
add_one / Jax / tpu / Primal 0.0000014628 s 0.000001404425 s 1.04
add_one / HLOOpt / tpu / Primal 0.00000145215 s 0.0000014319249999999998 s 1.01
add_one / PartOpt / tpu / Primal 0.0000014492749999999998 s 0.000001407725 s 1.03
add_one / IPartOpt / tpu / Primal 0.0000014420250000000002 s 0.0000014238 s 1.01
add_one / DefOpt / tpu / Primal 0.0000014542 s 0.0000014118 s 1.03
add_one / IDefOpt / tpu / Primal 0.0000014516000000000002 s 0.0000014299 s 1.02
add_one / JaXPipe / tpu / Forward 0.0000019020500000000005 s 0.000001843975 s 1.03
add_one / Jax / tpu / Forward 0.000001867375 s 0.000001847625 s 1.01
add_one / HLOOpt / tpu / Forward 0.0000019126 s 0.000001858425 s 1.03
add_one / PartOpt / tpu / Forward 0.000001863575 s 0.00000184495 s 1.01
add_one / IPartOpt / tpu / Forward 0.000001908975 s 0.000001852075 s 1.03
add_one / DefOpt / tpu / Forward 0.0000018687 s 0.000001841675 s 1.01
add_one / IDefOpt / tpu / Forward 0.000001918825 s 0.00000186845 s 1.03
add_one / JaXPipe / tpu / PreRev 0.00000226595 s 0.00000224345 s 1.01
add_one / JaXPipe / tpu / PostRev 0.0000022933999999999995 s 0.00000225025 s 1.02
add_one / JaXPipe / tpu / BothRev 0.000002264 s 0.000002233775 s 1.01
add_one / Jax / tpu / BothRev 0.000002302325 s 0.000002238575 s 1.03
add_one / HLOOpt / tpu / PreRev 0.000002255975 s 0.00000223345 s 1.01
add_one / HLOOpt / tpu / PostRev 0.000002295775 s 0.000002241375 s 1.02
add_one / HLOOpt / tpu / BothRev 0.0000022627 s 0.000002236875 s 1.01
add_one / PartOpt / tpu / PreRev 0.00000230295 s 0.00000223495 s 1.03
add_one / PartOpt / tpu / PostRev 0.00000226325 s 0.000002233225 s 1.01
add_one / PartOpt / tpu / BothRev 0.000002299425 s 0.000002241225 s 1.03
add_one / IPartOpt / tpu / PreRev 0.0000022572500000000003 s 0.000002242525 s 1.01
add_one / IPartOpt / tpu / PostRev 0.00000230195 s 0.00000223895 s 1.03
add_one / IPartOpt / tpu / BothRev 0.000002266875 s 0.000002238 s 1.01
add_one / DefOpt / tpu / PreRev 0.000002295975 s 0.0000022502 s 1.02
add_one / DefOpt / tpu / PostRev 0.000002262575 s 0.000002244275 s 1.01
add_one / DefOpt / tpu / BothRev 0.0000023021 s 0.0000022360000000000003 s 1.03
add_one / IDefOpt / tpu / PreRev 0.000002252225 s 0.000002237875 s 1.01
add_one / IDefOpt / tpu / PostRev 0.0000022949 s 0.000002241175 s 1.02
add_one / IDefOpt / tpu / BothRev 0.000002254425 s 0.000002246625 s 1.00
add_one / JaXPipe / cpu / Primal 0.000016467 s 0.000006647580121352803 s 2.48
add_one / Jax / cpu / Primal 0.000015838 s 0.000007019260065135313 s 2.26
add_one / HLOOpt / cpu / Primal 0.000015935000000000002 s 0.000006855379924672889 s 2.32
add_one / PartOpt / cpu / Primal 0.000015931 s 0.000006245519980438985 s 2.55
add_one / IPartOpt / cpu / Primal 0.000015918000000000002 s 0.000006942599975445774 s 2.29
add_one / DefOpt / cpu / Primal 0.000016154 s 0.0000069157200050540265 s 2.34
add_one / IDefOpt / cpu / Primal 0.00002129 s 0.00000644126001134282 s 3.31
add_one / JaXPipe / cpu / Forward 0.000022176 s 0.000010058619955088945 s 2.20
add_one / Jax / cpu / Forward 0.000022055 s 0.000009927860028255965 s 2.22
add_one / HLOOpt / cpu / Forward 0.000021967 s 0.000010169480028707766 s 2.16
add_one / PartOpt / cpu / Forward 0.000021781 s 0.000010235980043944435 s 2.13
add_one / IPartOpt / cpu / Forward 0.00002195 s 0.00000981865994617692 s 2.24
add_one / DefOpt / cpu / Forward 0.000022314 s 0.000010122280018549644 s 2.20
add_one / IDefOpt / cpu / Forward 0.000021381 s 0.000009593060021870769 s 2.23
add_one / JaXPipe / cpu / PreRev 0.000024594 s 0.000011487320007290693 s 2.14
add_one / JaXPipe / cpu / PostRev 0.000024101 s 0.000011399719842302149 s 2.11
add_one / JaXPipe / cpu / BothRev 0.000024394 s 0.000011612500002229353 s 2.10
add_one / Jax / cpu / BothRev 0.000024083 s 0.000011088899955211672 s 2.17
add_one / HLOOpt / cpu / PreRev 0.000024067 s 0.000011683019965857968 s 2.06
add_one / HLOOpt / cpu / PostRev 0.000029907 s 0.000014739599992026342 s 2.03
add_one / HLOOpt / cpu / BothRev 0.000023861 s 0.000011134360011055832 s 2.14
add_one / PartOpt / cpu / PreRev 0.000024295 s 0.0000115139000081399 s 2.11
add_one / PartOpt / cpu / PostRev 0.000024783 s 0.000011024920004274463 s 2.25
add_one / PartOpt / cpu / BothRev 0.000024077 s 0.00001160763993539149 s 2.07
add_one / IPartOpt / cpu / PreRev 0.000024317 s 0.000011384020108380355 s 2.14
add_one / IPartOpt / cpu / PostRev 0.000024346 s 0.000011333639959048014 s 2.15
add_one / IPartOpt / cpu / BothRev 0.000023822 s 0.000010866900011023972 s 2.19
add_one / DefOpt / cpu / PreRev 0.000024354 s 0.00001100651994420332 s 2.21
add_one / DefOpt / cpu / PostRev 0.00002442 s 0.000010951919994113267 s 2.23
add_one / DefOpt / cpu / BothRev 0.00002431 s 0.000010872719994949876 s 2.24
add_one / IDefOpt / cpu / PreRev 0.000024284 s 0.000011335459948895732 s 2.14
add_one / IDefOpt / cpu / PostRev 0.000024523 s 0.000010873259907384635 s 2.26
add_one / IDefOpt / cpu / BothRev 0.000024505 s 0.000010865979947993764 s 2.26
add_two / JaXPipe / cpu / Primal 0.000006840200001079211 s 0.000006799959974159719 s 1.01
add_two / Jax / cpu / Primal 0.00000754888000074061 s 0.000006684760064672446 s 1.13
add_two / HLOOpt / cpu / Primal 0.000007194499992237979 s 0.000006848320026620058 s 1.05
add_two / PartOpt / cpu / Primal 0.000006502599999294034 s 0.000007211239972093608 s 0.90
add_two / IPartOpt / cpu / Primal 0.000006943439998394751 s 0.0000069097999767109285 s 1.00
add_two / DefOpt / cpu / Primal 0.000006740120006725192 s 0.000007047919862088747 s 0.96
add_two / IDefOpt / cpu / Primal 0.000006646180004281632 s 0.0000069972399978723845 s 0.95
add_two / JaXPipe / cpu / Forward 0.000009834519999003532 s 0.00000989261992799584 s 0.99
add_two / Jax / cpu / Forward 0.00001018796000380462 s 0.000009906680061249062 s 1.03
add_two / HLOOpt / cpu / Forward 0.000010495540007013916 s 0.000010638600033416878 s 0.99
add_two / PartOpt / cpu / Forward 0.00001061793999042493 s 0.000010299079913238528 s 1.03
add_two / IPartOpt / cpu / Forward 0.000010410239995053416 s 0.000010367220074840588 s 1.00
add_two / DefOpt / cpu / Forward 0.000010335479996683717 s 0.000010392959993623662 s 0.99
add_two / IDefOpt / cpu / Forward 0.000010201379996033211 s 0.000010211160042672416 s 1.00
add_two / JaXPipe / cpu / PreRev 0.00001414407999618561 s 0.000013334040031622865 s 1.06
add_two / JaXPipe / cpu / PostRev 0.000014027679997070663 s 0.000013018559948250183 s 1.08
add_two / JaXPipe / cpu / BothRev 0.00001359044000537324 s 0.00001362989989502239 s 1.00
add_two / Jax / cpu / BothRev 0.000014596580019770044 s 0.00001345006005067262 s 1.09
add_two / HLOOpt / cpu / PreRev 0.000013970800007427896 s 0.000013947060087957652 s 1.00
add_two / HLOOpt / cpu / PostRev 0.00001607333999572802 s 0.000015459519945579815 s 1.04
add_two / HLOOpt / cpu / BothRev 0.00001372699999365068 s 0.000013334459890756988 s 1.03
add_two / PartOpt / cpu / PreRev 0.000013794439983030316 s 0.00001374192004732322 s 1.00
add_two / PartOpt / cpu / PostRev 0.00001399025999717196 s 0.000013714780016016448 s 1.02
add_two / PartOpt / cpu / BothRev 0.000013962059988443798 s 0.000014078620006330312 s 0.99
add_two / IPartOpt / cpu / PreRev 0.000014025279986071835 s 0.000013926319952588527 s 1.01
add_two / IPartOpt / cpu / PostRev 0.000013653699991209578 s 0.000013750919915764824 s 0.99
add_two / IPartOpt / cpu / BothRev 0.000014107999993484554 s 0.000012942820067110006 s 1.09
add_two / DefOpt / cpu / PreRev 0.000014240280011108553 s 0.000013346799914870643 s 1.07
add_two / DefOpt / cpu / PostRev 0.000013722239998514851 s 0.000013584420030383624 s 1.01
add_two / DefOpt / cpu / BothRev 0.000013793740004075517 s 0.000013460420032060938 s 1.02
add_two / IDefOpt / cpu / PreRev 0.000013849500001015256 s 0.000013556180074374424 s 1.02
add_two / IDefOpt / cpu / PostRev 0.00001361916000405472 s 0.0000135386000147264 s 1.01
add_two / IDefOpt / cpu / BothRev 0.000014261980002174823 s 0.0000137441399601812 s 1.04
add_two / JaXPipe / cuda / Primal 0.0000019200000000000003 s 0.000002432 s 0.79
add_two / Jax / cuda / Primal 0.0000019200000000000003 s 0.000002431 s 0.79
add_two / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.000002431 s 0.79
add_two / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.000002432 s 0.79
add_two / IPartOpt / cuda / Primal 0.000001919 s 0.000002431 s 0.79
add_two / DefOpt / cuda / Primal 0.000001919 s 0.000002431 s 0.79
add_two / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.000002432 s 0.79
add_two / JaXPipe / cuda / Forward 0.000009952 s 0.000010688 s 0.93
add_two / Jax / cuda / Forward 0.000009952 s 0.0000104 s 0.96
add_two / HLOOpt / cuda / Forward 0.00001008 s 0.000010432 s 0.97
add_two / PartOpt / cuda / Forward 0.000010016 s 0.000010432 s 0.96
add_two / IPartOpt / cuda / Forward 0.000009792 s 0.000010432 s 0.94
add_two / DefOpt / cuda / Forward 0.000010112 s 0.000010495 s 0.96
add_two / IDefOpt / cuda / Forward 0.0000096 s 0.000010432 s 0.92
add_two / JaXPipe / cuda / PreRev 0.000033792000000000004 s 0.000031904000000000005 s 1.06
add_two / JaXPipe / cuda / PostRev 0.000032767999999999995 s 0.000032767999999999995 s 1
add_two / JaXPipe / cuda / BothRev 0.000033984 s 0.000032256 s 1.05
add_two / Jax / cuda / BothRev 0.000033503 s 0.000032704 s 1.02
add_two / HLOOpt / cuda / PreRev 0.000033503 s 0.000032800000000000004 s 1.02
add_two / HLOOpt / cuda / PostRev 0.000033247 s 0.000032288 s 1.03
add_two / HLOOpt / cuda / BothRev 0.000032704 s 0.000033344 s 0.98
add_two / PartOpt / cuda / PreRev 0.000033376 s 0.000033344 s 1.00
add_two / PartOpt / cuda / PostRev 0.000032512 s 0.00003248 s 1.00
add_two / PartOpt / cuda / BothRev 0.000033568 s 0.00003248 s 1.03
add_two / IPartOpt / cuda / PreRev 0.000033598999999999995 s 0.000033152000000000004 s 1.01
add_two / IPartOpt / cuda / PostRev 0.000032896000000000005 s 0.000032992 s 1.00
add_two / IPartOpt / cuda / BothRev 0.000033632 s 0.000032416 s 1.04
add_two / DefOpt / cuda / PreRev 0.000033824 s 0.000032896000000000005 s 1.03
add_two / DefOpt / cuda / PostRev 0.000033695 s 0.000032543 s 1.04
add_two / DefOpt / cuda / BothRev 0.0000336 s 0.00003296 s 1.02
add_two / IDefOpt / cuda / PreRev 0.000034368 s 0.000033472 s 1.03
add_two / IDefOpt / cuda / PostRev 0.00003344 s 0.00003264 s 1.02
add_two / IDefOpt / cuda / BothRev 0.000033249 s 0.000033472 s 0.99
add_two / JaXPipe / tpu / Primal 0.0000014071749999999998 s 0.0000014749499999999995 s 0.95
add_two / Jax / tpu / Primal 0.0000014406749999999998 s 0.000001439825 s 1.00
add_two / HLOOpt / tpu / Primal 0.0000013922000000000002 s 0.0000014717749999999998 s 0.95
add_two / PartOpt / tpu / Primal 0.0000014516000000000002 s 0.000001445775 s 1.00
add_two / IPartOpt / tpu / Primal 0.00000140515 s 0.00000146895 s 0.96
add_two / DefOpt / tpu / Primal 0.000001461775 s 0.0000014387 s 1.02
add_two / IDefOpt / tpu / Primal 0.0000013948250000000002 s 0.00000147485 s 0.95
add_two / JaXPipe / tpu / Forward 0.000001798625 s 0.00000181685 s 0.99
add_two / Jax / tpu / Forward 0.0000017979 s 0.000001905375 s 0.94
add_two / HLOOpt / tpu / Forward 0.000001796 s 0.000001819725 s 0.99
add_two / PartOpt / tpu / Forward 0.0000017853749999999998 s 0.00000192145 s 0.93
add_two / IPartOpt / tpu / Forward 0.0000017973 s 0.000001815275 s 0.99
add_two / DefOpt / tpu / Forward 0.000001790475 s 0.00000190725 s 0.94
add_two / IDefOpt / tpu / Forward 0.000001808175 s 0.000001819325 s 0.99
add_two / JaXPipe / tpu / PreRev 0.0000028107 s 0.000002864375 s 0.98
add_two / JaXPipe / tpu / PostRev 0.00000272635 s 0.0000027442250000000004 s 0.99
add_two / JaXPipe / tpu / BothRev 0.00000279145 s 0.0000028623 s 0.98
add_two / Jax / tpu / BothRev 0.000002717775 s 0.000002725525 s 1.00
add_two / HLOOpt / tpu / PreRev 0.000002794975 s 0.0000028695500000000003 s 0.97
add_two / HLOOpt / tpu / PostRev 0.000002721825 s 0.0000027291 s 1.00
add_two / HLOOpt / tpu / BothRev 0.0000027924500000000003 s 0.0000028571 s 0.98
add_two / PartOpt / tpu / PreRev 0.000002728075 s 0.0000027297500000000004 s 1.00
add_two / PartOpt / tpu / PostRev 0.000002791825 s 0.000002863175 s 0.98
add_two / PartOpt / tpu / BothRev 0.00000271665 s 0.0000027407 s 0.99
add_two / IPartOpt / tpu / PreRev 0.0000028034750000000003 s 0.000002863175 s 0.98
add_two / IPartOpt / tpu / PostRev 0.000002721825 s 0.0000027239 s 1.00
add_two / IPartOpt / tpu / BothRev 0.0000028091000000000003 s 0.00000286345 s 0.98
add_two / DefOpt / tpu / PreRev 0.000002726375 s 0.0000027285749999999995 s 1.00
add_two / DefOpt / tpu / PostRev 0.0000028034750000000003 s 0.000002867725 s 0.98
add_two / DefOpt / tpu / BothRev 0.0000027312 s 0.00000272025 s 1.00
add_two / IDefOpt / tpu / PreRev 0.0000028021 s 0.000002858225 s 0.98
add_two / IDefOpt / tpu / PostRev 0.000002730525 s 0.00000273385 s 1.00
add_two / IDefOpt / tpu / BothRev 0.000002797325 s 0.0000028672 s 0.98
add_two / JaXPipe / cpu / Primal 0.000016735 s 0.000006799959974159719 s 2.46
add_two / Jax / cpu / Primal 0.000016726999999999998 s 0.000006684760064672446 s 2.50
add_two / HLOOpt / cpu / Primal 0.00001669 s 0.000006848320026620058 s 2.44
add_two / PartOpt / cpu / Primal 0.000016416 s 0.000007211239972093608 s 2.28
add_two / IPartOpt / cpu / Primal 0.000016632 s 0.0000069097999767109285 s 2.41
add_two / DefOpt / cpu / Primal 0.000016947 s 0.000007047919862088747 s 2.40
add_two / IDefOpt / cpu / Primal 0.000016376 s 0.0000069972399978723845 s 2.34
add_two / JaXPipe / cpu / Forward 0.00002257 s 0.00000989261992799584 s 2.28
add_two / Jax / cpu / Forward 0.000022361 s 0.000009906680061249062 s 2.26
add_two / HLOOpt / cpu / Forward 0.000021912 s 0.000010638600033416878 s 2.06
add_two / PartOpt / cpu / Forward 0.000022223 s 0.000010299079913238528 s 2.16
add_two / IPartOpt / cpu / Forward 0.000022495 s 0.000010367220074840588 s 2.17
add_two / DefOpt / cpu / Forward 0.000021924 s 0.000010392959993623662 s 2.11
add_two / IDefOpt / cpu / Forward 0.000022515 s 0.000010211160042672416 s 2.20
add_two / JaXPipe / cpu / PreRev 0.000028492 s 0.000013334040031622865 s 2.14
add_two / JaXPipe / cpu / PostRev 0.00002804 s 0.000013018559948250183 s 2.15
add_two / JaXPipe / cpu / BothRev 0.000028118 s 0.00001362989989502239 s 2.06
add_two / Jax / cpu / BothRev 0.000027874 s 0.00001345006005067262 s 2.07
add_two / HLOOpt / cpu / PreRev 0.000028483 s 0.000013947060087957652 s 2.04
add_two / HLOOpt / cpu / PostRev 0.000028493 s 0.000015459519945579815 s 1.84
add_two / HLOOpt / cpu / BothRev 0.000028444 s 0.000013334459890756988 s 2.13
add_two / PartOpt / cpu / PreRev 0.000028188 s 0.00001374192004732322 s 2.05
add_two / PartOpt / cpu / PostRev 0.00002871 s 0.000013714780016016448 s 2.09
add_two / PartOpt / cpu / BothRev 0.000028366 s 0.000014078620006330312 s 2.01
add_two / IPartOpt / cpu / PreRev 0.000028918 s 0.000013926319952588527 s 2.08
add_two / IPartOpt / cpu / PostRev 0.000028357 s 0.000013750919915764824 s 2.06
add_two / IPartOpt / cpu / BothRev 0.000028079 s 0.000012942820067110006 s 2.17
add_two / DefOpt / cpu / PreRev 0.00002881 s 0.000013346799914870643 s 2.16
add_two / DefOpt / cpu / PostRev 0.000028608 s 0.000013584420030383624 s 2.11
add_two / DefOpt / cpu / BothRev 0.000028282 s 0.000013460420032060938 s 2.10
add_two / IDefOpt / cpu / PreRev 0.000028423 s 0.000013556180074374424 s 2.10
add_two / IDefOpt / cpu / PostRev 0.000028279 s 0.0000135386000147264 s 2.09
add_two / IDefOpt / cpu / BothRev 0.000028329 s 0.0000137441399601812 s 2.06
cache / JaXPipe / cpu / Primal 0.000006446539987337019 s 0.000006199420040502446 s 1.04
cache / Jax / cpu / Primal 0.0000066530800108921535 s 0.000006489560000773053 s 1.03
cache / HLOOpt / cpu / Primal 0.00000625936000233196 s 0.0000061433001064870044 s 1.02
cache / PartOpt / cpu / Primal 0.000006805019991134031 s 0.0000059924800007138405 s 1.14
cache / IPartOpt / cpu / Primal 0.000006479320004473265 s 0.000006026699920766987 s 1.08
cache / DefOpt / cpu / Primal 0.000006537159997606068 s 0.000006346639984258218 s 1.03
cache / IDefOpt / cpu / Primal 0.000006523500014736783 s 0.000006013300044287462 s 1.08
cache / JaXPipe / cpu / Forward 0.000013845040007254285 s 0.000014293200001702644 s 0.97
cache / Jax / cpu / Forward 0.00001370008000094458 s 0.000018218660061393164 s 0.75
cache / HLOOpt / cpu / Forward 0.00001486359999034903 s 0.000015171059985732426 s 0.98
cache / PartOpt / cpu / Forward 0.000014036860002306638 s 0.000014781880017835648 s 0.95
cache / IPartOpt / cpu / Forward 0.0000143534400012868 s 0.000015402959925268077 s 0.93
cache / DefOpt / cpu / Forward 0.000014431880001666288 s 0.000015190559988695895 s 0.95
cache / IDefOpt / cpu / Forward 0.000014161699991745991 s 0.000015089259977685288 s 0.94
cache / JaXPipe / cpu / PreRev 0.000014620500003275083 s 0.00001638670000829734 s 0.89
cache / JaXPipe / cpu / PostRev 0.00001921986000297693 s 0.00002091621987347025 s 0.92
cache / JaXPipe / cpu / BothRev 0.000015062980010043248 s 0.0000167550800506433 s 0.90
cache / Jax / cpu / BothRev 0.00001949208000269209 s 0.000020225059906806563 s 0.96
cache / HLOOpt / cpu / PreRev 0.000016221860000769084 s 0.000016973180026980118 s 0.96
cache / HLOOpt / cpu / PostRev 0.00001875815999937913 s 0.00001933609994011931 s 0.97
cache / HLOOpt / cpu / BothRev 0.00001658344000361467 s 0.00001763584001309937 s 0.94
cache / PartOpt / cpu / PreRev 0.0000146189800011598 s 0.00001631077997444663 s 0.90
cache / PartOpt / cpu / PostRev 0.00002016495999896506 s 0.00002172737995351781 s 0.93
cache / PartOpt / cpu / BothRev 0.000015417440001783687 s 0.000016777699966041836 s 0.92
cache / IPartOpt / cpu / PreRev 0.000015057279988468508 s 0.000016093579997686902 s 0.94
cache / IPartOpt / cpu / PostRev 0.000019447279985342902 s 0.00002099820003422792 s 0.93
cache / IPartOpt / cpu / BothRev 0.000014837260007425357 s 0.000016570839961786986 s 0.90
cache / DefOpt / cpu / PreRev 0.000015220939997107052 s 0.00001723014000162948 s 0.88
cache / DefOpt / cpu / PostRev 0.000014355060002344545 s 0.000017786020052881212 s 0.81
cache / DefOpt / cpu / BothRev 0.000015280920003988285 s 0.00001701712002613931 s 0.90
cache / IDefOpt / cpu / PreRev 0.00001582859999871289 s 0.000016709280007489724 s 0.95
cache / IDefOpt / cpu / PostRev 0.00001538348000622136 s 0.000017414860012650023 s 0.88
cache / IDefOpt / cpu / BothRev 0.000014632679997248488 s 0.000015970800013747067 s 0.92
cache / JaXPipe / cuda / Primal 0.000002304 s 0.000002335 s 0.99
cache / Jax / cuda / Primal 0.000002335 s 0.000002335 s 1
cache / HLOOpt / cuda / Primal 0.000002272 s 0.000002335 s 0.97
cache / PartOpt / cuda / Primal 0.000002272 s 0.000002335 s 0.97
cache / IPartOpt / cuda / Primal 0.000002335 s 0.000002335 s 1
cache / DefOpt / cuda / Primal 0.000002272 s 0.000002335 s 0.97
cache / IDefOpt / cuda / Primal 0.00000224 s 0.000002335 s 0.96
cache / JaXPipe / cuda / Forward 0.000002336 s 0.0000023670000000000004 s 0.99
cache / Jax / cuda / Forward 0.000002335 s 0.000002336 s 1.00
cache / HLOOpt / cuda / Forward 0.000002335 s 0.0000023670000000000004 s 0.99
cache / PartOpt / cuda / Forward 0.000002335 s 0.0000023670000000000004 s 0.99
cache / IPartOpt / cuda / Forward 0.000002336 s 0.0000023670000000000004 s 0.99
cache / DefOpt / cuda / Forward 0.000002271 s 0.0000023670000000000004 s 0.96
cache / IDefOpt / cuda / Forward 0.000002335 s 0.0000023670000000000004 s 0.99
cache / JaXPipe / cuda / PreRev 0.00001248 s 0.000011231 s 1.11
cache / JaXPipe / cuda / PostRev 0.000010816 s 0.000010944 s 0.99
cache / JaXPipe / cuda / BothRev 0.000010944 s 0.00001072 s 1.02
cache / Jax / cuda / BothRev 0.000010944 s 0.000010687 s 1.02
cache / HLOOpt / cuda / PreRev 0.000013376 s 0.000013696 s 0.98
cache / HLOOpt / cuda / PostRev 0.00001344 s 0.0000136 s 0.99
cache / HLOOpt / cuda / BothRev 0.000013375 s 0.000013663 s 0.98
cache / PartOpt / cuda / PreRev 0.000011008 s 0.00001072 s 1.03
cache / PartOpt / cuda / PostRev 0.00001088 s 0.000010688 s 1.02
cache / PartOpt / cuda / BothRev 0.000011296 s 0.000010784 s 1.05
cache / IPartOpt / cuda / PreRev 0.000010944 s 0.000010816 s 1.01
cache / IPartOpt / cuda / PostRev 0.000010816 s 0.000010656 s 1.02
cache / IPartOpt / cuda / BothRev 0.000012448 s 0.000010879 s 1.14
cache / DefOpt / cuda / PreRev 0.000010944 s 0.000011615 s 0.94
cache / DefOpt / cuda / PostRev 0.000011071 s 0.000010528 s 1.05
cache / DefOpt / cuda / BothRev 0.00001088 s 0.000010592 s 1.03
cache / IDefOpt / cuda / PreRev 0.000010816 s 0.000010624 s 1.02
cache / IDefOpt / cuda / PostRev 0.000010432 s 0.000010912 s 0.96
cache / IDefOpt / cuda / BothRev 0.000010751 s 0.000010752 s 1.00
cache / JaXPipe / tpu / Primal 0.0000024557 s 0.000002460175 s 1.00
cache / Jax / tpu / Primal 0.00000246525 s 0.00000246785 s 1.00
cache / HLOOpt / tpu / Primal 0.0000024685750000000003 s 0.00000248385 s 0.99
cache / PartOpt / tpu / Primal 0.00000246285 s 0.0000024633 s 1.00
cache / IPartOpt / tpu / Primal 0.000002471025 s 0.000002465975 s 1.00
cache / DefOpt / tpu / Primal 0.000002470375 s 0.000002450575 s 1.01
cache / IDefOpt / tpu / Primal 0.000002453875 s 0.0000024649 s 1.00
cache / JaXPipe / tpu / Forward 0.0000035457 s 0.000003556525 s 1.00
cache / Jax / tpu / Forward 0.0000035270000000000003 s 0.000003554925 s 0.99
cache / HLOOpt / tpu / Forward 0.0000035413 s 0.000003565225 s 0.99
cache / PartOpt / tpu / Forward 0.0000035469 s 0.00000356175 s 1.00
cache / IPartOpt / tpu / Forward 0.000003549875 s 0.00000357025 s 0.99
cache / DefOpt / tpu / Forward 0.00000353635 s 0.000003556825 s 0.99
cache / IDefOpt / tpu / Forward 0.0000035463 s 0.000003585375 s 0.99
cache / JaXPipe / tpu / PreRev 0.00000493095 s 0.000004975425000000001 s 0.99
cache / JaXPipe / tpu / PostRev 0.00000499015 s 0.000004970025 s 1.00
cache / JaXPipe / tpu / BothRev 0.000004992750000000001 s 0.000004990425 s 1.00
cache / Jax / tpu / BothRev 0.000004991475 s 0.000004974374999999999 s 1.00
cache / HLOOpt / tpu / PreRev 0.000004120725 s 0.0000039449 s 1.04
cache / HLOOpt / tpu / PostRev 0.0000041521 s 0.0000041285 s 1.01
cache / HLOOpt / tpu / BothRev 0.0000041272 s 0.00000395235 s 1.04
cache / PartOpt / tpu / PreRev 0.0000049899 s 0.000004977275 s 1.00
cache / PartOpt / tpu / PostRev 0.00000497465 s 0.0000049768750000000005 s 1.00
cache / PartOpt / tpu / BothRev 0.000004986225000000001 s 0.000004991225 s 1.00
cache / IPartOpt / tpu / PreRev 0.00000495895 s 0.00000500825 s 0.99
cache / IPartOpt / tpu / PostRev 0.000005001475 s 0.00000496675 s 1.01
cache / IPartOpt / tpu / BothRev 0.000004971324999999999 s 0.0000049771 s 1.00
cache / DefOpt / tpu / PreRev 0.000005003300000000001 s 0.0000049878 s 1.00
cache / DefOpt / tpu / PostRev 0.000004971 s 0.0000049536 s 1.00
cache / DefOpt / tpu / BothRev 0.000004989275000000001 s 0.000004968224999999999 s 1.00
cache / IDefOpt / tpu / PreRev 0.0000049629 s 0.0000049741 s 1.00
cache / IDefOpt / tpu / PostRev 0.000004996275 s 0.000004959374999999999 s 1.01
cache / IDefOpt / tpu / BothRev 0.0000049668250000000005 s 0.0000049691500000000005 s 1.00
cache / JaXPipe / cpu / Primal 0.000018451 s 0.000006199420040502446 s 2.98
cache / Jax / cpu / Primal 0.000018338 s 0.000006489560000773053 s 2.83
cache / HLOOpt / cpu / Primal 0.000024816 s 0.0000061433001064870044 s 4.04
cache / PartOpt / cpu / Primal 0.00001813 s 0.0000059924800007138405 s 3.03
cache / IPartOpt / cpu / Primal 0.000018191 s 0.000006026699920766987 s 3.02
cache / DefOpt / cpu / Primal 0.000018614 s 0.000006346639984258218 s 2.93
cache / IDefOpt / cpu / Primal 0.000018094 s 0.000006013300044287462 s 3.01
cache / JaXPipe / cpu / Forward 0.000020988 s 0.000014293200001702644 s 1.47
cache / Jax / cpu / Forward 0.00002184 s 0.000018218660061393164 s 1.20
cache / HLOOpt / cpu / Forward 0.00002145 s 0.000015171059985732426 s 1.41
cache / PartOpt / cpu / Forward 0.000021744 s 0.000014781880017835648 s 1.47
cache / IPartOpt / cpu / Forward 0.0000211 s 0.000015402959925268077 s 1.37
cache / DefOpt / cpu / Forward 0.000021509 s 0.000015190559988695895 s 1.42
cache / IDefOpt / cpu / Forward 0.000023517 s 0.000015089259977685288 s 1.56
cache / JaXPipe / cpu / PreRev 0.000023159 s 0.00001638670000829734 s 1.41
cache / JaXPipe / cpu / PostRev 0.000026558 s 0.00002091621987347025 s 1.27
cache / JaXPipe / cpu / BothRev 0.0000223 s 0.0000167550800506433 s 1.33
cache / Jax / cpu / BothRev 0.000025062 s 0.000020225059906806563 s 1.24
cache / HLOOpt / cpu / PreRev 0.000022532 s 0.000016973180026980118 s 1.33
cache / HLOOpt / cpu / PostRev 0.000021247 s 0.00001933609994011931 s 1.10
cache / HLOOpt / cpu / BothRev 0.000022104 s 0.00001763584001309937 s 1.25
cache / PartOpt / cpu / PreRev 0.000021889 s 0.00001631077997444663 s 1.34
cache / PartOpt / cpu / PostRev 0.00002733 s 0.00002172737995351781 s 1.26
cache / PartOpt / cpu / BothRev 0.000021621 s 0.000016777699966041836 s 1.29
cache / IPartOpt / cpu / PreRev 0.00002132 s 0.000016093579997686902 s 1.32
cache / IPartOpt / cpu / PostRev 0.00003021 s 0.00002099820003422792 s 1.44
cache / IPartOpt / cpu / BothRev 0.000022012 s 0.000016570839961786986 s 1.33
cache / DefOpt / cpu / PreRev 0.000022158 s 0.00001723014000162948 s 1.29
cache / DefOpt / cpu / PostRev 0.000022118 s 0.000017786020052881212 s 1.24
cache / DefOpt / cpu / BothRev 0.000021693 s 0.00001701712002613931 s 1.27
cache / IDefOpt / cpu / PreRev 0.000022039 s 0.000016709280007489724 s 1.32
cache / IDefOpt / cpu / PostRev 0.000022225 s 0.000017414860012650023 s 1.28
cache / IDefOpt / cpu / BothRev 0.000022168 s 0.000015970800013747067 s 1.39
Concat / JaXPipe / cpu / Primal 0.000006987860001572699 s 0.000006435699961002684 s 1.09
Concat / Jax / cpu / Primal 0.000007094180014064477 s 0.000006501879979623481 s 1.09
Concat / HLOOpt / cpu / Primal 0.000007173079993663123 s 0.000006443000111175934 s 1.11
Concat / PartOpt / cpu / Primal 0.000006423200011340668 s 0.00000656270001854864 s 0.98
Concat / IPartOpt / cpu / Primal 0.000006700220003494906 s 0.000006713340062560746 s 1.00
Concat / DefOpt / cpu / Primal 0.000006941719991573336 s 0.000006703879953420255 s 1.04
Concat / IDefOpt / cpu / Primal 0.0000067667399980564365 s 0.000006299739998212317 s 1.07
Concat / JaXPipe / cpu / Forward 0.000009941479995632108 s 0.000009630539989302633 s 1.03
Concat / Jax / cpu / Forward 0.000010117180006545822 s 0.000010102839951287024 s 1.00
Concat / HLOOpt / cpu / Forward 0.000009940780003034889 s 0.000009807360020204217 s 1.01
Concat / PartOpt / cpu / Forward 0.0000103309199926116 s 0.000009883120037557092 s 1.05
Concat / IPartOpt / cpu / Forward 0.000009755899993706407 s 0.00001014442004816374 s 0.96
Concat / DefOpt / cpu / Forward 0.000009747140015861078 s 0.000009600979974493384 s 1.02
Concat / IDefOpt / cpu / Forward 0.000010225700007140404 s 0.000009490459960943554 s 1.08
Concat / JaXPipe / cpu / PreRev 0.000012085939990811312 s 0.00001145785996413906 s 1.05
Concat / JaXPipe / cpu / PostRev 0.000012182500011022058 s 0.00001061990007656277 s 1.15
Concat / JaXPipe / cpu / BothRev 0.000011586059993078378 s 0.00001069060004738276 s 1.08
Concat / Jax / cpu / BothRev 0.000011706680011229764 s 0.00001097654005207005 s 1.07
Concat / HLOOpt / cpu / PreRev 0.000012278600004265171 s 0.000011447080014477251 s 1.07
Concat / HLOOpt / cpu / PostRev 0.00001411556000221026 s 0.000013339120014279616 s 1.06
Concat / HLOOpt / cpu / BothRev 0.00001124364001043432 s 0.000011410279857955175 s 0.99
Concat / PartOpt / cpu / PreRev 0.00001187844000014593 s 0.000010918520038103452 s 1.09
Concat / PartOpt / cpu / PostRev 0.000011328659991249878 s 0.000011513339977682337 s 0.98
Concat / PartOpt / cpu / BothRev 0.000011523639998358705 s 0.000011210639950149926 s 1.03
Concat / IPartOpt / cpu / PreRev 0.000012632059999759805 s 0.000011497539999254512 s 1.10
Concat / IPartOpt / cpu / PostRev 0.000011799219996646571 s 0.00001105878000089433 s 1.07
Concat / IPartOpt / cpu / BothRev 0.000011713439998857212 s 0.000011240539970458486 s 1.04
Concat / DefOpt / cpu / PreRev 0.000012313140000514978 s 0.000011512620076246096 s 1.07
Concat / DefOpt / cpu / PostRev 0.000012430500014488643 s 0.00001148805995399016 s 1.08
Concat / DefOpt / cpu / BothRev 0.000011841260006804075 s 0.000010975040077028096 s 1.08
Concat / IDefOpt / cpu / PreRev 0.00001208582000572278 s 0.000011794299989560388 s 1.02
Concat / IDefOpt / cpu / PostRev 0.000011327040006108291 s 0.00001099843995689298 s 1.03
Concat / IDefOpt / cpu / BothRev 0.000012178239994682372 s 0.00001133895997554646 s 1.07
Concat / JaXPipe / cuda / Primal 0.000001919 s 0.000002463 s 0.78
Concat / Jax / cuda / Primal 0.0000019200000000000003 s 0.000002463 s 0.78
Concat / HLOOpt / cuda / Primal 0.0000019200000000000003 s 0.000002463 s 0.78
Concat / PartOpt / cuda / Primal 0.0000019200000000000003 s 0.000002463 s 0.78
Concat / IPartOpt / cuda / Primal 0.0000019200000000000003 s 0.000002463 s 0.78
Concat / DefOpt / cuda / Primal 0.0000019200000000000003 s 0.000002464 s 0.78
Concat / IDefOpt / cuda / Primal 0.0000019200000000000003 s 0.000002463 s 0.78
Concat / JaXPipe / cuda / Forward 0.00000992 s 0.000010688 s 0.93
Concat / Jax / cuda / Forward 0.000010015 s 0.000010592 s 0.95
Concat / HLOOpt / cuda / Forward 0.000011168 s 0.000010752 s 1.04
Concat / PartOpt / cuda / Forward 0.00000992 s 0.000010592 s 0.94
Concat / IPartOpt / cuda / Forward 0.00001008 s 0.000010528 s 0.96
Concat / DefOpt / cuda / Forward 0.000010176 s 0.000010688 s 0.95
Concat / IDefOpt / cuda / Forward 0.00001024 s 0.00001072 s 0.96
Concat / JaXPipe / cuda / PreRev 0.000016544 s 0.000017216 s 0.96
Concat / JaXPipe / cuda / PostRev 0.00001696 s 0.000016896000000000002 s 1.00
Concat / JaXPipe / cuda / BothRev 0.000016608 s 0.000017152 s 0.97
Concat / Jax / cuda / BothRev 0.000016670999999999997 s 0.000017088 s 0.98
Concat / HLOOpt / cuda / PreRev 0.000016768000000000003 s 0.000016864 s 0.99
Concat / HLOOpt / cuda / PostRev 0.000016832 s 0.00001728 s 0.97
Concat / HLOOpt / cuda / BothRev 0.000016864 s 0.000017247999999999998 s 0.98
Concat / PartOpt / cuda / PreRev 0.000016992 s 0.000017088 s 0.99
Concat / PartOpt / cuda / PostRev 0.000016416 s 0.000017215 s 0.95
Concat / PartOpt / cuda / BothRev 0.000016607 s 0.000017152 s 0.97
Concat / IPartOpt / cuda / PreRev 0.00001712 s 0.000017024 s 1.01
Concat / IPartOpt / cuda / PostRev 0.000017024 s 0.000016768000000000003 s 1.02
Concat / IPartOpt / cuda / BothRev 0.00001664 s 0.000016768999999999998 s 0.99
Concat / DefOpt / cuda / PreRev 0.000016512 s 0.00001728 s 0.96
Concat / DefOpt / cuda / PostRev 0.000016927999999999998 s 0.00001712 s 0.99
Concat / DefOpt / cuda / BothRev 0.000016864 s 0.000016927999999999998 s 1.00
Concat / IDefOpt / cuda / PreRev 0.000016863 s 0.0000176 s 0.96
Concat / IDefOpt / cuda / PostRev 0.000016288 s 0.000016831 s 0.97
Concat / IDefOpt / cuda / BothRev 0.000016607 s 0.000017152 s 0.97
Concat / JaXPipe / tpu / Primal 0.0000015222 s 0.00000152845 s 1.00
Concat / Jax / tpu / Primal 0.0000015192249999999998 s 0.000001536025 s 0.99
Concat / HLOOpt / tpu / Primal 0.0000015191750000000002 s 0.000001532375 s 0.99
Concat / PartOpt / tpu / Primal 0.000001520025 s 0.0000015328249999999998 s 0.99
Concat / IPartOpt / tpu / Primal 0.0000015187 s 0.00000152755 s 0.99
Concat / DefOpt / tpu / Primal 0.00000151085 s 0.000001544 s 0.98
Concat / IDefOpt / tpu / Primal 0.000001525725 s 0.00000153295 s 1.00
Concat / JaXPipe / tpu / Forward 0.0000015486 s 0.000001575425 s 0.98
Concat / Jax / tpu / Forward 0.000001562425 s 0.0000015991 s 0.98
Concat / HLOOpt / tpu / Forward 0.00000153035 s 0.00000158675 s 0.96
Concat / PartOpt / tpu / Forward 0.00000154955 s 0.0000015972500000000005 s 0.97
Concat / IPartOpt / tpu / Forward 0.000001555425 s 0.000001587375 s 0.98
Concat / DefOpt / tpu / Forward 0.0000015516250000000002 s 0.00000161155 s 0.96
Concat / IDefOpt / tpu / Forward 0.0000015611249999999998 s 0.0000015929 s 0.98
Concat / JaXPipe / tpu / PreRev 0.0000020263 s 0.000002013 s 1.01
Concat / JaXPipe / tpu / PostRev 0.000001997375 s 0.0000020712 s 0.96
Concat / JaXPipe / tpu / BothRev 0.00000202705 s 0.00000201835 s 1.00
Concat / Jax / tpu / BothRev 0.00000199485 s 0.000002055325 s 0.97
Concat / HLOOpt / tpu / PreRev 0.00000202175 s 0.000002012675 s 1.00
Concat / HLOOpt / tpu / PostRev 0.00000200035 s 0.000002068325 s 0.97
Concat / HLOOpt / tpu / BothRev 0.0000020276 s 0.000002007525 s 1.01
Concat / PartOpt / tpu / PreRev 0.0000019986 s 0.000002057425 s 0.97
Concat / PartOpt / tpu / PostRev 0.00000202085 s 0.0000020110000000000003 s 1.00
Concat / PartOpt / tpu / BothRev 0.000001996725 s 0.0000020543000000000004 s 0.97
Concat / IPartOpt / tpu / PreRev 0.0000020210250000000003 s 0.00000201685 s 1.00
Concat / IPartOpt / tpu / PostRev 0.000001994775 s 0.0000020674 s 0.96
Concat / IPartOpt / tpu / BothRev 0.0000020260500000000003 s 0.00000201065 s 1.01
Concat / DefOpt / tpu / PreRev 0.0000019972750000000003 s 0.000002056325 s 0.97
Concat / DefOpt / tpu / PostRev 0.0000020210250000000003 s 0.0000020085 s 1.01
Concat / DefOpt / tpu / BothRev 0.0000019922500000000003 s 0.000002056375 s 0.97
Concat / IDefOpt / tpu / PreRev 0.0000020248 s 0.000002020325 s 1.00
Concat / IDefOpt / tpu / PostRev 0.0000019928250000000003 s 0.000002062825 s 0.97
Concat / IDefOpt / tpu / BothRev 0.0000020228 s 0.000002019675 s 1.00
Concat / JaXPipe / cpu / Primal 0.000016018 s 0.000006435699961002684 s 2.49
Concat / Jax / cpu / Primal 0.000015890999999999997 s 0.000006501879979623481 s 2.44
Concat / HLOOpt / cpu / Primal 0.000015762999999999998 s 0.000006443000111175934 s 2.45
Concat / PartOpt / cpu / Primal 0.000015585 s 0.00000656270001854864 s 2.37
Concat / IPartOpt / cpu / Primal 0.000016018 s 0.000006713340062560746 s 2.39
Concat / DefOpt / cpu / Primal 0.000015857 s 0.000006703879953420255 s 2.37
Concat / IDefOpt / cpu / Primal 0.000015978 s 0.000006299739998212317 s 2.54
Concat / JaXPipe / cpu / Forward 0.000022134 s 0.000009630539989302633 s 2.30
Concat / Jax / cpu / Forward 0.000021513 s 0.000010102839951287024 s 2.13
Concat / HLOOpt / cpu / Forward 0.000021723 s 0.000009807360020204217 s 2.21
Concat / PartOpt / cpu / Forward 0.000021453 s 0.000009883120037557092 s 2.17
Concat / IPartOpt / cpu / Forward 0.000021729 s 0.00001014442004816374 s 2.14
Concat / DefOpt / cpu / Forward 0.000021859 s 0.000009600979974493384 s 2.28
Concat / IDefOpt / cpu / Forward 0.000021772 s 0.000009490459960943554 s 2.29
Concat / JaXPipe / cpu / PreRev 0.000025244 s 0.00001145785996413906 s 2.20
Concat / JaXPipe / cpu / PostRev 0.000024598 s 0.00001061990007656277 s 2.32
Concat / JaXPipe / cpu / BothRev 0.000024477 s 0.00001069060004738276 s 2.29
Concat / Jax / cpu / BothRev 0.000024925 s 0.00001097654005207005 s 2.27
Concat / HLOOpt / cpu / PreRev 0.000030452 s 0.000011447080014477251 s 2.66
Concat / HLOOpt / cpu / PostRev 0.000024378 s 0.000013339120014279616 s 1.83
Concat / HLOOpt / cpu / BothRev 0.000024272 s 0.000011410279857955175 s 2.13
Concat / PartOpt / cpu / PreRev 0.000024412 s 0.000010918520038103452 s 2.24
Concat / PartOpt / cpu / PostRev 0.000024086 s 0.000011513339977682337 s 2.09
Concat / PartOpt / cpu / BothRev 0.000023779 s 0.000011210639950149926 s 2.12
Concat / IPartOpt / cpu / PreRev 0.000024947 s 0.000011497539999254512 s 2.17
Concat / IPartOpt / cpu / PostRev 0.000024024 s 0.00001105878000089433 s 2.17
Concat / IPartOpt / cpu / BothRev 0.00002421 s 0.000011240539970458486 s 2.15
Concat / DefOpt / cpu / PreRev 0.000024505 s 0.000011512620076246096 s 2.13
Concat / DefOpt / cpu / PostRev 0.00002441 s 0.00001148805995399016 s 2.12
Concat / DefOpt / cpu / BothRev 0.000024428 s 0.000010975040077028096 s 2.23
Concat / IDefOpt / cpu / PreRev 0.000024647 s 0.000011794299989560388 s 2.09
Concat / IDefOpt / cpu / PostRev 0.000024088 s 0.00001099843995689298 s 2.19
Concat / IDefOpt / cpu / BothRev 0.000024431 s 0.00001133895997554646 s 2.15
const_scatter / JaXPipe / cpu / Primal 0.000006625020000683435 s 0.000006160580014693551 s 1.08
const_scatter / Jax / cpu / Primal 0.000006743919996097247 s 0.0000061382400235743264 s 1.10
const_scatter / HLOOpt / cpu / Primal 0.000007086320003963919 s 0.000007276999967871234 s 0.97
const_scatter / PartOpt / cpu / Primal 0.000006670980008038896 s 0.000006049840012565256 s 1.10
const_scatter / IPartOpt / cpu / Primal 0.000006939020006484498 s 0.000006259759975364432 s 1.11
const_scatter / DefOpt / cpu / Primal 0.000006778400002076524 s 0.000007348239942075452 s 0.92
const_scatter / IDefOpt / cpu / Primal 0.000006903020005211147 s 0.00000687206003931351 s 1.00
const_scatter / JaXPipe / cpu / Forward 0.000011372240003311162 s 0.000010107159941981082 s 1.13
const_scatter / Jax / cpu / Forward 0.000009195459992952238 s 0.000008885880015441217 s 1.03
const_scatter / HLOOpt / cpu / Forward 0.000011250579989336984 s 0.000010267419947922462 s 1.10
const_scatter / PartOpt / cpu / Forward 0.000010812000000441911 s 0.0000099613599559234 s 1.09
const_scatter / IPartOpt / cpu / Forward 0.000010957480003526144 s 0.000010300219983037096 s 1.06
const_scatter / DefOpt / cpu / Forward 0.000010814380004831037 s 0.000009794319976208498 s 1.10
const_scatter / IDefOpt / cpu / Forward 0.000010805919994254508 s 0.00001040247994751553 s 1.04
const_scatter / JaXPipe / cpu / PreRev 0.0002898265599992 s 0.0002957961599895 s 0.98
const_scatter / JaXPipe / cpu / PostRev 0.0002815211599931 s 0.0002822270399155 s 1.00
const_scatter / JaXPipe / cpu / BothRev 0.0002978083200036 s 0.0002834643999995 s 1.05
const_scatter / Jax / cpu / BothRev 0.0002852807999806 s 0.0002829939000366 s 1.01
const_scatter / HLOOpt / cpu / PreRev 0.0002829849800036 s 0.0002844214600008 s 0.99
const_scatter / HLOOpt / cpu / PostRev 0.0002863282200064 s 0.0002877555599661 s 1.00
const_scatter / HLOOpt / cpu / BothRev 0.0002858254200009 s 0.0002840586799902 s 1.01
const_scatter / PartOpt / cpu / PreRev 0.0002831315600019 s 0.0002818456199747 s 1.00
const_scatter / PartOpt / cpu / PostRev 0.0002990480599964 s 0.0002807304599446 s 1.07
const_scatter / PartOpt / cpu / BothRev 0.0002835402999949 s 0.0002830758599884 s 1.00
const_scatter / IPartOpt / cpu / PreRev 0.000284828579995 s 0.000283077799977 s 1.01
const_scatter / IPartOpt / cpu / PostRev 0.0002796286600096 s 0.0002849672000593 s 0.98
const_scatter / IPartOpt / cpu / BothRev 0.0002813879399946 s 0.0002838028199948 s 0.99
const_scatter / DefOpt / cpu / PreRev 0.0002827757399973 s 0.0002814923199548 s 1.00
const_scatter / DefOpt / cpu / PostRev 0.0002828244800116 s 0.0002810001200305 s 1.01
const_scatter / DefOpt / cpu / BothRev 0.0002820929599943 s 0.000283605760087 s 0.99
const_scatter / IDefOpt / cpu / PreRev 0.0002822563400104 s 0.0002903664798395 s 0.97
const_scatter / IDefOpt / cpu / PostRev 0.0002842051600077 s 0.0002862955600539 s 0.99
const_scatter / IDefOpt / cpu / BothRev 0.0002851890800047 s 0.0002827165999951 s 1.01
const_scatter / JaXPipe / cuda / Primal 0.000001887 s 0.000002463 s 0.77
const_scatter / Jax / cuda / Primal 0.000001887 s 0.000002464 s 0.77
const_scatter / HLOOpt / cuda / Primal 0.000001887 s 0.000002463 s 0.77
const_scatter / PartOpt / cuda / Primal 0.000001887 s 0.000002464 s 0.77
const_scatter / IPartOpt / cuda / Primal 0.000001887 s 0.000002463 s 0.77
const_scatter / DefOpt / cuda / Primal 0.000001887 s 0.000002463 s 0.77
const_scatter / IDefOpt / cuda / Primal 0.000001887 s 0.000002463 s 0.77
const_scatter / JaXPipe / cuda / Forward 0.000009536 s 0.00001056 s 0.90
const_scatter / Jax / cuda / Forward 0.000010176 s 0.00001056 s 0.96
const_scatter / HLOOpt / cuda / Forward 0.000010048 s 0.000010496 s 0.96
const_scatter / PartOpt / cuda / Forward 0.000010304 s 0.00001056 s 0.98
const_scatter / IPartOpt / cuda / Forward 0.000009344 s 0.000010624 s 0.88
const_scatter / DefOpt / cuda / Forward 0.000010176 s 0.000010687 s 0.95
const_scatter / IDefOpt / cuda / Forward 0.000010304 s 0.000010879 s 0.95
const_scatter / JaXPipe / cuda / PreRev 0.000016576000000000002 s 0.00001664 s 1.00
const_scatter / JaXPipe / cuda / PostRev 0.000016927999999999998 s 0.000017184 s 0.99
const_scatter / JaXPipe / cuda / BothRev 0.000016896000000000002 s 0.000017375999999999998 s 0.97
const_scatter / Jax / cuda / BothRev 0.000016864 s 0.000017695 s 0.95
const_scatter / HLOOpt / cuda / PreRev 0.000016607 s 0.000016513 s 1.01
const_scatter / HLOOpt / cuda / PostRev 0.000017184 s 0.00001696 s 1.01
const_scatter / HLOOpt / cuda / BothRev 0.0000168 s 0.000017375000000000002 s 0.97
const_scatter / PartOpt / cuda / PreRev 0.000016927999999999998 s 0.000017375999999999998 s 0.97
const_scatter / PartOpt / cuda / PostRev 0.00001664 s 0.0000168 s 0.99
const_scatter / PartOpt / cuda / BothRev 0.00001664 s 0.000017152 s 0.97
const_scatter / IPartOpt / cuda / PreRev 0.000016544 s 0.000017056 s 0.97
const_scatter / IPartOpt / cuda / PostRev 0.000017247999999999998 s 0.000016927999999999998 s 1.02
const_scatter / IPartOpt / cuda / BothRev 0.00001664 s 0.000016416 s 1.01
const_scatter / DefOpt / cuda / PreRev 0.000016576000000000002 s 0.000017183 s 0.96
const_scatter / DefOpt / cuda / PostRev 0.000016511 s 0.000016736 s 0.99
const_scatter / DefOpt / cuda / BothRev 0.000016447 s 0.000016768000000000003 s 0.98
const_scatter / IDefOpt / cuda / PreRev 0.000016512 s 0.000016224 s 1.02
const_scatter / IDefOpt / cuda / PostRev 0.000016768000000000003 s 0.000016575 s 1.01
const_scatter / IDefOpt / cuda / BothRev 0.000016576000000000002 s 0.000016864 s 0.98
const_scatter / JaXPipe / tpu / Primal 0.00000381455 s 0.00000379245 s 1.01
const_scatter / Jax / tpu / Primal 0.000003848525 s 0.000003796425 s 1.01
const_scatter / HLOOpt / tpu / Primal 0.000003781075 s 0.000003792775 s 1.00
const_scatter / PartOpt / tpu / Primal 0.000003834125 s 0.00000380505 s 1.01
const_scatter / IPartOpt / tpu / Primal 0.00000381335 s 0.00000382235 s 1.00
const_scatter / DefOpt / tpu / Primal 0.0000038215 s 0.000003791025 s 1.01
const_scatter / IDefOpt / tpu / Primal 0.0000037844 s 0.000003789375 s 1.00
const_scatter / JaXPipe / tpu / Forward 0.00000650605 s 0.0000064895 s 1.00
const_scatter / Jax / tpu / Forward 0.00000648185 s 0.00000645865 s 1.00
const_scatter / HLOOpt / tpu / Forward 0.000006483974999999999 s 0.000006481249999999999 s 1.00
const_scatter / PartOpt / tpu / Forward 0.000006469675 s 0.000006462725 s 1.00
const_scatter / IPartOpt / tpu / Forward 0.000006505699999999999 s 0.000006488175000000001 s 1.00
const_scatter / DefOpt / tpu / Forward 0.000006452774999999999 s 0.0000064513 s 1.00
const_scatter / IDefOpt / tpu / Forward 0.000006495524999999999 s 0.000006512999999999999 s 1.00
const_scatter / JaXPipe / tpu / PreRev 0.000006639575 s 0.00000668195 s 0.99
const_scatter / JaXPipe / tpu / PostRev 0.0000066269 s 0.000006673599999999999 s 0.99
const_scatter / JaXPipe / tpu / BothRev 0.000006631725 s 0.000006703925 s 0.99
const_scatter / Jax / tpu / BothRev 0.0000066388 s 0.0000066635 s 1.00
const_scatter / HLOOpt / tpu / PreRev 0.000006614275 s 0.00000667115 s 0.99
const_scatter / HLOOpt / tpu / PostRev 0.00000665485 s 0.000006664799999999999 s 1.00
const_scatter / HLOOpt / tpu / BothRev 0.000006583575000000001 s 0.000006675999999999999 s 0.99
const_scatter / PartOpt / tpu / PreRev 0.0000066486 s 0.0000066796750000000006 s 1.00
const_scatter / PartOpt / tpu / PostRev 0.000006588825 s 0.0000066963 s 0.98
const_scatter / PartOpt / tpu / BothRev 0.000006648575 s 0.0000066842 s 0.99
const_scatter / IPartOpt / tpu / PreRev 0.000006612775 s 0.000006679924999999999 s 0.99
const_scatter / IPartOpt / tpu / PostRev 0.0000066399 s 0.00000668675 s 0.99
const_scatter / IPartOpt / tpu / BothRev 0.000006618725 s 0.0000066709 s 0.99
const_scatter / DefOpt / tpu / PreRev 0.000006634425 s 0.00000668525 s 0.99
const_scatter / DefOpt / tpu / PostRev 0.000006614974999999999 s 0.000006678425 s 0.99
const_scatter / DefOpt / tpu / BothRev 0.0000066429 s 0.00000669215 s 0.99
const_scatter / IDefOpt / tpu / PreRev 0.000006600125 s 0.0000066717 s 0.99
const_scatter / IDefOpt / tpu / PostRev 0.000006641850000000001 s 0.00000667215 s 1.00
const_scatter / IDefOpt / tpu / BothRev 0.000006605125 s 0.00000666985 s 0.99
const_scatter / JaXPipe / cpu / Primal 0.000016202999999999997 s 0.000006160580014693551 s 2.63
const_scatter / Jax / cpu / Primal 0.000015662 s 0.0000061382400235743264 s 2.55
const_scatter / HLOOpt / cpu / Primal 0.000016503 s 0.000007276999967871234 s 2.27
const_scatter / PartOpt / cpu / Primal 0.000015577 s 0.000006049840012565256 s 2.57
const_scatter / IPartOpt / cpu / Primal 0.000015713 s 0.000006259759975364432 s 2.51
const_scatter / DefOpt / cpu / Primal 0.000016575 s 0.000007348239942075452 s 2.26
const_scatter / IDefOpt / cpu / Primal 0.000016871 s 0.00000687206003931351 s 2.46
const_scatter / JaXPipe / cpu / Forward 0.000022505 s 0.000010107159941981082 s 2.23
const_scatter / Jax / cpu / Forward 0.000020568 s 0.000008885880015441217 s 2.31
const_scatter / HLOOpt / cpu / Forward 0.000022429 s 0.000010267419947922462 s 2.18
const_scatter / PartOpt / cpu / Forward 0.000022245 s 0.0000099613599559234 s 2.23
const_scatter / IPartOpt / cpu / Forward 0.00002249 s 0.000010300219983037096 s 2.18
const_scatter / DefOpt / cpu / Forward 0.000022087 s 0.000009794319976208498 s 2.26
const_scatter / IDefOpt / cpu / Forward 0.000022413 s 0.00001040247994751553 s 2.15
const_scatter / JaXPipe / cpu / PreRev 0.000530271 s 0.0002957961599895 s 1.79
const_scatter / JaXPipe / cpu / PostRev 0.00051888 s 0.0002822270399155 s 1.84
const_scatter / JaXPipe / cpu / BothRev 0.000538464 s 0.0002834643999995 s 1.90
const_scatter / Jax / cpu / BothRev 0.000530301 s 0.0002829939000366 s 1.87
const_scatter / HLOOpt / cpu / PreRev 0.000518402 s 0.0002844214600008 s 1.82
const_scatter / HLOOpt / cpu / PostRev 0.000528954 s 0.0002877555599661 s 1.84
const_scatter / HLOOpt / cpu / BothRev 0.000528486 s 0.0002840586799902 s 1.86
const_scatter / PartOpt / cpu / PreRev 0.0005278469999999 s 0.0002818456199747 s 1.87
const_scatter / PartOpt / cpu / PostRev 0.000541846 s 0.0002807304599446 s 1.93
const_scatter / PartOpt / cpu / BothRev 0.000530961 s 0.0002830758599884 s 1.88
const_scatter / IPartOpt / cpu / PreRev 0.000515733 s 0.000283077799977 s 1.82
const_scatter / IPartOpt / cpu / PostRev 0.0005402 s 0.0002849672000593 s 1.90
const_scatter / IPartOpt / cpu / BothRev 0.000529247 s 0.0002838028199948 s 1.86
const_scatter / DefOpt / cpu / PreRev 0.00051905 s 0.0002814923199548 s 1.84
const_scatter / DefOpt / cpu / PostRev 0.0005208779999999 s 0.0002810001200305 s 1.85
const_scatter / DefOpt / cpu / BothRev 0.000530159 s 0.000283605760087 s 1.87
const_scatter / IDefOpt / cpu / PreRev 0.00053315 s 0.0002903664798395 s 1.84
const_scatter / IDefOpt / cpu / PostRev 0.000523737 s 0.0002862955600539 s 1.83
const_scatter / IDefOpt / cpu / BothRev 0.000529992 s 0.0002827165999951 s 1.87
GenDot / JaXPipe / cpu / Primal 0.00000686431999611159 s 0.000007377620004263008 s 0.93
GenDot / Jax / cpu / Primal 0.000006981219992212573 s 0.000007222200038086157 s 0.97
GenDot / HLOOpt / cpu / Primal 0.000007963359998939267 s 0.000007422399994538864 s 1.07
GenDot / PartOpt / cpu / Primal 0.0000067217000014352376 s 0.000006913960078236414 s 0.97
GenDot / IPartOpt / cpu / Primal 0.000006923240000560327 s 0.0000068490199737425424 s 1.01
GenDot / DefOpt / cpu / Primal 0.000007078400001319096 s 0.000007118580106180161 s 0.99
GenDot / IDefOpt / cpu / Primal 0.000006958339995435381 s 0.000007142399899748853 s 0.97
GenDot / JaXPipe / cpu / Forward 0.000010823920001712395 s 0.000010665599984349682 s 1.01
GenDot / Jax / cpu / Forward 0.000010122919979949074 s 0.000010504599977139153 s 0.96
GenDot / HLOOpt / cpu / Forward 0.000011861380005484537 s 0.000011046119925595122 s 1.07
GenDot / PartOpt / cpu / Forward 0.000010922419996859389 s 0.00001006010003038682 s 1.09
GenDot / IPartOpt / cpu / Forward 0.000011176879995673517 s 0.00001120744003856089 s 1.00
GenDot / DefOpt / cpu / Forward 0.000011151740011428046 s 0.000010274079904775135 s 1.09
GenDot / IDefOpt / cpu / Forward 0.000010538600010931986 s 0.000010486480096005835 s 1.00
GenDot / JaXPipe / cpu / PreRev 0.000011621360001754513 s 0.000010926600007223896 s 1.06
GenDot / JaXPipe / cpu / PostRev 0.0000106500799938658 s 0.00001013830002193572 s 1.05
GenDot / JaXPipe / cpu / BothRev 0.000010987260000092646 s 0.000010925779952231096 s 1.01
GenDot / Jax / cpu / BothRev 0.000010573260005912744 s 0.000009941319967765594 s 1.06
GenDot / HLOOpt / cpu / PreRev 0.00001126444000419724 s 0.000011691660110955128 s 0.96
GenDot / HLOOpt / cpu / PostRev 0.000013197619998663868 s 0.000012876239925390107 s 1.02
GenDot / HLOOpt / cpu / BothRev 0.000011241919994517957 s 0.00001098990001992206 s 1.02
GenDot / PartOpt / cpu / PreRev 0.000011715860007370792 s 0.000010831399958988184 s 1.08
GenDot / PartOpt / cpu / PostRev 0.000010246260001167683 s 0.000010444540021126158 s 0.98
GenDot / PartOpt / cpu / BothRev 0.000011765660001401555 s 0.000011348319931130391 s 1.04
GenDot / IPartOpt / cpu / PreRev 0.000011052820002532828 s 0.000010319640023226383 s 1.07
GenDot / IPartOpt / cpu / PostRev 0.000010592199996608544 s 0.000010269659978803248 s 1.03
GenDot / IPartOpt / cpu / BothRev 0.000010639479999099422 s 0.000011077979925175896 s 0.96
GenDot / DefOpt / cpu / PreRev 0.000010505199998078752 s 0.000010820759998750872 s 0.97
GenDot / DefOpt / cpu / PostRev 0.000011739920005311432 s 0.000010590620040602515 s 1.11
GenDot / DefOpt / cpu / BothRev 0.000010598139995181555 s 0.000010781660021166318 s 0.98
GenDot / IDefOpt / cpu / PreRev 0.000011468259999674049 s 0.00001114787997721578 s 1.03
GenDot / IDefOpt / cpu / PostRev 0.000011556360002487054 s 0.000011181120007677236 s 1.03
GenDot / IDefOpt / cpu / BothRev 0.000011499580004965535 s 0.000010971699994115623 s 1.05
GenDot / JaXPipe / cuda / Primal 0.000002015 s 0.000002528 s 0.80
GenDot / Jax / cuda / Primal 0.000002015 s 0.000002528 s 0.80
GenDot / HLOOpt / cuda / Primal 0.000001983 s 0.000002527 s 0.78
GenDot / PartOpt / cuda / Primal 0.000002015 s 0.00000256 s 0.79
GenDot / IPartOpt / cuda / Primal 0.000002015 s 0.00000256 s 0.79
GenDot / DefOpt / cuda / Primal 0.000001983 s 0.000002528 s 0.78
GenDot / IDefOpt / cuda / Primal 0.000001983 s 0.000002528 s 0.78
GenDot / JaXPipe / cuda / Forward 0.00001024 s 0.00001088 s 0.94
GenDot / Jax / cuda / Forward 0.000010208 s 0.000011776 s 0.87
GenDot / HLOOpt / cuda / Forward 0.000010176 s 0.000011872 s 0.86
GenDot / PartOpt / cuda / Forward 0.000010176 s 0.000010784 s 0.94
GenDot / IPartOpt / cuda / Forward 0.000010271 s 0.000010816 s 0.95
GenDot / DefOpt / cuda / Forward 0.000010016 s 0.000010656 s 0.94
GenDot / IDefOpt / cuda / Forward 0.000010016 s 0.000010624 s 0.94
GenDot / JaXPipe / cuda / PreRev 0.000009889 s 0.00001072 s 0.92
GenDot / JaXPipe / cuda / PostRev 0.000010047 s 0.000010656 s 0.94
GenDot / JaXPipe / cuda / BothRev 0.00001008 s 0.00001104 s 0.91
GenDot / Jax / cuda / BothRev 0.000010017 s 0.000011039 s 0.91
GenDot / HLOOpt / cuda / PreRev 0.000009824 s 0.000010753 s 0.91
GenDot / HLOOpt / cuda / PostRev 0.000010112 s 0.000010753 s 0.94
GenDot / HLOOpt / cuda / BothRev 0.000010112 s 0.000010752 s 0.94
GenDot / PartOpt / cuda / PreRev 0.000010336 s 0.000010912 s 0.95
GenDot / PartOpt / cuda / PostRev 0.000009728 s 0.000010816 s 0.90
GenDot / PartOpt / cuda / BothRev 0.000010496 s 0.000010752 s 0.98
GenDot / IPartOpt / cuda / PreRev 0.000010176 s 0.000010656 s 0.95
GenDot / IPartOpt / cuda / PostRev 0.000010624 s 0.000010912 s 0.97
GenDot / IPartOpt / cuda / BothRev 0.000010112 s 0.000010688 s 0.95
GenDot / DefOpt / cuda / PreRev 0.000011136 s 0.00001104 s 1.01
GenDot / DefOpt / cuda / PostRev 0.000010272 s 0.00001072 s 0.96
GenDot / DefOpt / cuda / BothRev 0.000011104 s 0.000010784 s 1.03
GenDot / IDefOpt / cuda / PreRev 0.000010335 s 0.000010783 s 0.96
GenDot / IDefOpt / cuda / PostRev 0.000010176 s 0.000010752 s 0.95
GenDot / IDefOpt / cuda / BothRev 0.000010144 s 0.000014816 s 0.68
GenDot / JaXPipe / tpu / Primal 9.431e-7 s 9.20625e-7 s 1.02
GenDot / Jax / tpu / Primal 9.2975e-7 s 9.299e-7 s 1.00
GenDot / HLOOpt / tpu / Primal 0.000001597975 s 0.0000016143 s 0.99
GenDot / PartOpt / tpu / Primal 9.3025e-7 s 9.30425e-7 s 1.00
GenDot / IPartOpt / tpu / Primal 9.434e-7 s 9.770499999999998e-7 s 0.97
GenDot / DefOpt / tpu / Primal 0.0000014993500000000003 s 0.000001497 s 1.00
GenDot / IDefOpt / tpu / Primal 0.00000159565 s 0.0000016151 s 0.99
GenDot / JaXPipe / tpu / Forward 0.0000030527 s 0.000003060275 s 1.00
GenDot / Jax / tpu / Forward 0.000002273475 s 0.00000231795 s 0.98
GenDot / HLOOpt / tpu / Forward 0.0000031117250000000003 s 0.0000031134 s 1.00
GenDot / PartOpt / tpu / Forward 0.0000031383 s 0.000003119825 s 1.01
GenDot / IPartOpt / tpu / Forward 0.0000031109750000000003 s 0.00000311875 s 1.00
GenDot / DefOpt / tpu / Forward 0.00000315315 s 0.000003119 s 1.01
GenDot / IDefOpt / tpu / Forward 0.00000310895 s 0.0000031241750000000005 s 1.00
GenDot / JaXPipe / tpu / PreRev 0.000003026575 s 0.000002942475 s 1.03
GenDot / JaXPipe / tpu / PostRev 0.0000023781 s 0.0000023475 s 1.01
GenDot / JaXPipe / tpu / BothRev 0.000003007275 s 0.000002937775 s 1.02
GenDot / Jax / tpu / BothRev 0.000002375875 s 0.00000235035 s 1.01
GenDot / HLOOpt / tpu / PreRev 0.0000030206 s 0.0000029276999999999995 s 1.03
GenDot / HLOOpt / tpu / PostRev 0.0000029344749999999995 s 0.000002871875 s 1.02
GenDot / HLOOpt / tpu / BothRev 0.000003024325 s 0.0000029336750000000003 s 1.03
GenDot / PartOpt / tpu / PreRev 0.000002928325 s 0.000002884175 s 1.02
GenDot / PartOpt / tpu / PostRev 0.000002413475 s 0.000002411675 s 1.00
GenDot / PartOpt / tpu / BothRev 0.000002938525 s 0.0000028810000000000005 s 1.02
GenDot / IPartOpt / tpu / PreRev 0.0000030098000000000004 s 0.0000029337 s 1.03
GenDot / IPartOpt / tpu / PostRev 0.0000023766499999999995 s 0.000002344475 s 1.01
GenDot / IPartOpt / tpu / BothRev 0.0000030002750000000003 s 0.00000292585 s 1.03
GenDot / DefOpt / tpu / PreRev 0.000002933075 s 0.000002879575 s 1.02
GenDot / DefOpt / tpu / PostRev 0.00000299965 s 0.00000294605 s 1.02
GenDot / DefOpt / tpu / BothRev 0.000002935125 s 0.00000288235 s 1.02
GenDot / IDefOpt / tpu / PreRev 0.00000301095 s 0.000002949175 s 1.02
GenDot / IDefOpt / tpu / PostRev 0.000002954725 s 0.0000028782 s 1.03
GenDot / IDefOpt / tpu / BothRev 0.0000030254500000000004 s 0.0000029348 s 1.03
GenDot / JaXPipe / cpu / Primal 0.000018086000000000003 s 0.000007377620004263008 s 2.45
GenDot / Jax / cpu / Primal 0.000018394 s 0.000007222200038086157 s 2.55
GenDot / HLOOpt / cpu / Primal 0.000017919999999999998 s 0.000007422399994538864 s 2.41
GenDot / PartOpt / cpu / Primal 0.000018124 s 0.000006913960078236414 s 2.62
GenDot / IPartOpt / cpu / Primal 0.000018707 s 0.0000068490199737425424 s 2.73
GenDot / DefOpt / cpu / Primal 0.000017107 s 0.000007118580106180161 s 2.40
GenDot / IDefOpt / cpu / Primal 0.0000174 s 0.000007142399899748853 s 2.44
GenDot / JaXPipe / cpu / Forward 0.000024053 s 0.000010665599984349682 s 2.26
GenDot / Jax / cpu / Forward 0.000024671 s 0.000010504599977139153 s 2.35
GenDot / HLOOpt / cpu / Forward 0.000024332 s 0.000011046119925595122 s 2.20
GenDot / PartOpt / cpu / Forward 0.000023896 s 0.00001006010003038682 s 2.38
GenDot / IPartOpt / cpu / Forward 0.000023904 s 0.00001120744003856089 s 2.13
GenDot / DefOpt / cpu / Forward 0.000023892 s 0.000010274079904775135 s 2.33
GenDot / IDefOpt / cpu / Forward 0.00002375 s 0.000010486480096005835 s 2.26
GenDot / JaXPipe / cpu / PreRev 0.000024306 s 0.000010926600007223896 s 2.22
GenDot / JaXPipe / cpu / PostRev 0.000042112 s 0.00001013830002193572 s 4.15
GenDot / JaXPipe / cpu / BothRev 0.000023767 s 0.000010925779952231096 s 2.18
GenDot / Jax / cpu / BothRev 0.00002525 s 0.000009941319967765594 s 2.54
GenDot / HLOOpt / cpu / PreRev 0.000029786 s 0.000011691660110955128 s 2.55
GenDot / HLOOpt / cpu / PostRev 0.000024401 s 0.000012876239925390107 s 1.90
GenDot / HLOOpt / cpu / BothRev 0.00002403 s 0.00001098990001992206 s 2.19
GenDot / PartOpt / cpu / PreRev 0.000023997 s 0.000010831399958988184 s 2.22
GenDot / PartOpt / cpu / PostRev 0.000025291 s 0.000010444540021126158 s 2.42
GenDot / PartOpt / cpu / BothRev 0.000024312 s 0.000011348319931130391 s 2.14
GenDot / IPartOpt / cpu / PreRev 0.000023949 s 0.000010319640023226383 s 2.32
GenDot / IPartOpt / cpu / PostRev 0.000025401 s 0.000010269659978803248 s 2.47
GenDot / IPartOpt / cpu / BothRev 0.000024251 s 0.000011077979925175896 s 2.19
GenDot / DefOpt / cpu / PreRev 0.000023925 s 0.000010820759998750872 s 2.21
GenDot / DefOpt / cpu / PostRev 0.000024429 s 0.000010590620040602515 s 2.31
GenDot / DefOpt / cpu / BothRev 0.00002434 s 0.000010781660021166318 s 2.26
GenDot / IDefOpt / cpu / PreRev 0.0000239 s 0.00001114787997721578 s 2.14
GenDot / IDefOpt / cpu / PostRev 0.000024272 s 0.000011181120007677236 s 2.17
GenDot / IDefOpt / cpu / BothRev 0.000024217 s 0.000010971699994115623 s 2.21
hlo_ffi / JaXPipe / cpu / Primal 0.00001110061999725076 s 0.00001053083999067894 s 1.05
hlo_ffi / Jax / cpu / Primal 0.000010119079993273772 s 0.000010049380052805646 s 1.01
hlo_ffi / HLOOpt / cpu / Primal 0.000010355979995892994 s 0.000010057840063382173 s 1.03
hlo_ffi / PartOpt / cpu / Primal 0.000009535500009860696 s 0.000009655160010879626 s 0.99
hlo_ffi / IPartOpt / cpu / Primal 0.00001084456000398859 s 0.00001009035997412866 s 1.07
hlo_ffi / DefOpt / cpu / Primal 0.000009700880007130765 s 0.000009545619977870956 s 1.02
hlo_ffi / IDefOpt / cpu / Primal 0.000009669100008977694 s 0.000010037020037998446 s 0.96
hlo_ffi / JaXPipe / cpu / Forward 0.000013769179993232683 s 0.00001457850004953798 s 0.94
hlo_ffi / Jax / cpu / Forward 0.000014144859997031744 s 0.000014279340030043386 s 0.99
hlo_ffi / HLOOpt / cpu / Forward 0.000014393319995633646 s 0.000014677200015285052 s 0.98
hlo_ffi / PartOpt / cpu / Forward 0.000013611200008654125 s 0.000014379120075318496 s 0.95
hlo_ffi / IPartOpt / cpu / Forward 0.00001414082000110284 s 0.000014442919964494649 s 0.98
hlo_ffi / DefOpt / cpu / Forward 0.000014447379992361676 s 0.00001432873994417605 s 1.01
hlo_ffi / IDefOpt / cpu / Forward 0.000014132520002476669 s 0.000013966859987704084 s 1.01
hlo_ffi / JaXPipe / cpu / PreRev 0.000014951280011246128 s 0.000014706940037285676 s 1.02
hlo_ffi / JaXPipe / cpu / PostRev 0.000013915320007527043 s 0.000014398339972103714 s 0.97
hlo_ffi / JaXPipe / cpu / BothRev 0.000014180379998833814 s 0.000014566539884981469 s 0.97
hlo_ffi / Jax / cpu / BothRev 0.000014619500004755536 s 0.000014863860051264056 s 0.98
hlo_ffi / HLOOpt / cpu / PreRev 0.000014733139998952538 s 0.00001461602003473672 s 1.01
hlo_ffi / HLOOpt / cpu / PostRev 0.00001625992000754195 s 0.000016265079939330462 s 1.00
hlo_ffi / HLOOpt / cpu / BothRev 0.000014234539994504305 s 0.000014199220004229574 s 1.00
hlo_ffi / PartOpt / cpu / PreRev 0.000014699799994559724 s 0.000014724999928148464 s 1.00
hlo_ffi / PartOpt / cpu / PostRev 0.000014566699994702504 s 0.00001431689992386964 s 1.02
hlo_ffi / PartOpt / cpu / BothRev 0.00001469587999054056 s 0.000013945200025773374 s 1.05
hlo_ffi / IPartOpt / cpu / PreRev 0.000014944859997285677 s 0.000014579059989046071 s 1.03
hlo_ffi / IPartOpt / cpu / PostRev 0.000014368620006735 s 0.000014473160026682308 s 0.99
hlo_ffi / IPartOpt / cpu / BothRev 0.000013914520015987364 s 0.000013991260002512718 s 0.99
hlo_ffi / DefOpt / cpu / PreRev 0.000014755160004824574 s 0.00001468752001528628 s 1.00
hlo_ffi / DefOpt / cpu / PostRev 0.000014391480010544913 s 0.000014148140035104009 s 1.02
hlo_ffi / DefOpt / cpu / BothRev 0.000013538220002828892 s 0.00001415964003172121 s 0.96
hlo_ffi / IDefOpt / cpu / PreRev 0.000014630799996666613 s 0.000014535159989463864 s 1.01
hlo_ffi / IDefOpt / cpu / PostRev 0.000014593840000998173 s 0.000014654960068583024 s 1.00
hlo_ffi / IDefOpt / cpu / BothRev 0.000014271679999637854 s 0.000014252880046115024 s 1.00
hlo_ffi / JaXPipe / cuda / Primal 0.000001983 s 0.000002368 s 0.84
hlo_ffi / Jax / cuda / Primal 0.000001984 s 0.000002368 s 0.84
hlo_ffi / HLOOpt / cuda / Primal 0.000001983 s 0.0000023670000000000004 s 0.84
hlo_ffi / PartOpt / cuda / Primal 0.000001983 s 0.000002368 s 0.84
hlo_ffi / IPartOpt / cuda / Primal 0.000001983 s 0.0000023670000000000004 s 0.84
hlo_ffi / DefOpt / cuda / Primal 0.000001983 s 0.0000023670000000000004 s 0.84
hlo_ffi / IDefOpt / cuda / Primal 0.000001983 s 0.0000023670000000000004 s 0.84
hlo_ffi / JaXPipe / cuda / Forward 0.00000208 s 0.000002463 s 0.84
hlo_ffi / Jax / cuda / Forward 0.000002048 s 0.000002463 s 0.83
hlo_ffi / HLOOpt / cuda / Forward 0.000002048 s 0.000002464 s 0.83
hlo_ffi / PartOpt / cuda / Forward 0.000002047 s 0.000002463 s 0.83
hlo_ffi / IPartOpt / cuda / Forward 0.000002048 s 0.000002463 s 0.83
hlo_ffi / DefOpt / cuda / Forward 0.000002048 s 0.000002463 s 0.83
hlo_ffi / IDefOpt / cuda / Forward 0.000002079 s 0.000002463 s 0.84
hlo_ffi / JaXPipe / cuda / PreRev 0.000002048 s 0.000002463 s 0.83
hlo_ffi / JaXPipe / cuda / PostRev 0.000002047 s 0.000002432 s 0.84
hlo_ffi / JaXPipe / cuda / BothRev 0.000002047 s 0.000002463 s 0.83
hlo_ffi / Jax / cuda / BothRev 0.000002048 s 0.000002432 s 0.84
hlo_ffi / HLOOpt / cuda / PreRev 0.000002047 s 0.000002463 s 0.83
hlo_ffi / HLOOpt / cuda / PostRev 0.000002048 s 0.000002463 s 0.83
hlo_ffi / HLOOpt / cuda / BothRev 0.000002047 s 0.000002432 s 0.84
hlo_ffi / PartOpt / cuda / PreRev 0.000002047 s 0.000002463 s 0.83
hlo_ffi / PartOpt / cuda / PostRev 0.000002048 s 0.000002463 s 0.83
hlo_ffi / PartOpt / cuda / BothRev 0.000002048 s 0.000002432 s 0.84
hlo_ffi / IPartOpt / cuda / PreRev 0.000002047 s 0.000002463 s 0.83
hlo_ffi / IPartOpt / cuda / PostRev 0.000002047 s 0.000002463 s 0.83
hlo_ffi / IPartOpt / cuda / BothRev 0.000002048 s 0.000002463 s 0.83
hlo_ffi / DefOpt / cuda / PreRev 0.000002047 s 0.000002463 s 0.83
hlo_ffi / DefOpt / cuda / PostRev 0.000002047 s 0.000002433 s 0.84
hlo_ffi / DefOpt / cuda / BothRev 0.000002048 s 0.000002432 s 0.84
hlo_ffi / IDefOpt / cuda / PreRev 0.000002047 s 0.000002432 s 0.84
hlo_ffi / IDefOpt / cuda / PostRev 0.000002048 s 0.000002432 s 0.84
hlo_ffi / IDefOpt / cuda / BothRev 0.000002047 s 0.000002433 s 0.84
hlo_ffi / JaXPipe / tpu / Primal 9.24025e-7 s 9.3165e-7 s 0.99
hlo_ffi / Jax / tpu / Primal 9.51875e-7 s 9.53075e-7 s 1.00
hlo_ffi / HLOOpt / tpu / Primal 8.952000000000001e-7 s 9.0595e-7 s 0.99
hlo_ffi / PartOpt / tpu / Primal 9.53175e-7 s 9.50075e-7 s 1.00
hlo_ffi / IPartOpt / tpu / Primal 8.971999999999999e-7 s 9.0665e-7 s 0.99
hlo_ffi / DefOpt / tpu / Primal 9.50275e-7 s 9.52375e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Primal 8.972500000000001e-7 s 9.084e-7 s 0.99
hlo_ffi / JaXPipe / tpu / Forward 9.49175e-7 s 9.49675e-7 s 1.00
hlo_ffi / Jax / tpu / Forward 9.81925e-7 s 9.81725e-7 s 1.00
hlo_ffi / HLOOpt / tpu / Forward 9.74075e-7 s 9.74525e-7 s 1.00
hlo_ffi / PartOpt / tpu / Forward 9.3415e-7 s 9.344e-7 s 1.00
hlo_ffi / IPartOpt / tpu / Forward 9.7415e-7 s 9.74575e-7 s 1.00
hlo_ffi / DefOpt / tpu / Forward 9.344e-7 s 9.345e-7 s 1.00
hlo_ffi / IDefOpt / tpu / Forward 9.74325e-7 s 9.736749999999998e-7 s 1.00
hlo_ffi / JaXPipe / tpu / PreRev 9.32425e-7 s 9.38e-7 s 0.99
hlo_ffi / JaXPipe / tpu / PostRev 9.65575e-7 s 9.656e-7 s 1.00
hlo_ffi / JaXPipe / tpu / BothRev 9.625e-7 s 9.6225e-7 s 1.00
hlo_ffi / Jax / tpu / BothRev 9.6505e-7 s 9.6555e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PreRev 9.62325e-7 s 9.6295e-7 s 1.00
hlo_ffi / HLOOpt / tpu / PostRev 9.64975e-7 s 9.65025e-7 s 1.00
hlo_ffi / HLOOpt / tpu / BothRev 9.6205e-7 s 9.62275e-7 s 1.00
hlo_ffi / PartOpt / tpu / PreRev 9.652e-7 s 9.65375e-7 s 1.00
hlo_ffi / PartOpt / tpu / PostRev 9.618500000000002e-7 s 9.615750000000002e-7 s 1.00
hlo_ffi / PartOpt / tpu / BothRev 9.64475e-7 s 9.6505e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PreRev 9.616e-7 s 9.62725e-7 s 1.00
hlo_ffi / IPartOpt / tpu / PostRev 9.654e-7 s 9.64875e-7 s 1.00
hlo_ffi / IPartOpt / tpu / BothRev 9.61825e-7 s 9.62275e-7 s 1.00
hlo_ffi / DefOpt / tpu / PreRev 9.6515e-7 s 9.646e-7 s 1.00
hlo_ffi / DefOpt / tpu / PostRev 9.619e-7 s 9.62475e-7 s 1.00
hlo_ffi / DefOpt / tpu / BothRev 9.6495e-7 s 9.64825e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PreRev 9.617e-7 s 9.625749999999998e-7 s 1.00
hlo_ffi / IDefOpt / tpu / PostRev 9.6555e-7 s 9.65e-7 s 1.00
hlo_ffi / IDefOpt / tpu / BothRev 9.62125e-7 s 9.62625e-7 s 1.00
hlo_ffi / JaXPipe / cpu / Primal 0.000022439 s 0.00001053083999067894 s 2.13
hlo_ffi / Jax / cpu / Primal 0.000022134 s 0.000010049380052805646 s 2.20
hlo_ffi / HLOOpt / cpu / Primal 0.000022387 s 0.000010057840063382173 s 2.23
hlo_ffi / PartOpt / cpu / Primal 0.00002187 s 0.000009655160010879626 s 2.27
hlo_ffi / IPartOpt / cpu / Primal 0.000021709 s 0.00001009035997412866 s 2.15
hlo_ffi / DefOpt / cpu / Primal 0.000021823 s 0.000009545619977870956 s 2.29
hlo_ffi / IDefOpt / cpu / Primal 0.000021924 s 0.000010037020037998446 s 2.18
hlo_ffi / JaXPipe / cpu / Forward 0.000030439 s 0.00001457850004953798 s 2.09
hlo_ffi / Jax / cpu / Forward 0.000029874 s 0.000014279340030043386 s 2.09
hlo_ffi / HLOOpt / cpu / Forward 0.000029823 s 0.000014677200015285052 s 2.03
hlo_ffi / PartOpt / cpu / Forward 0.000030014 s 0.000014379120075318496 s 2.09
hlo_ffi / IPartOpt / cpu / Forward 0.000029983 s 0.000014442919964494649 s 2.08
hlo_ffi / DefOpt / cpu / Forward 0.000029885 s 0.00001432873994417605 s 2.09
hlo_ffi / IDefOpt / cpu / Forward 0.00002972 s 0.000013966859987704084 s 2.13
hlo_ffi / JaXPipe / cpu / PreRev 0.000030206 s 0.000014706940037285676 s 2.05
hlo_ffi / JaXPipe / cpu / PostRev 0.000029307 s 0.000014398339972103714 s 2.04
hlo_ffi / JaXPipe / cpu / BothRev 0.000029197 s 0.000014566539884981469 s 2.00
hlo_ffi / Jax / cpu / BothRev 0.000029608 s 0.000014863860051264056 s 1.99
hlo_ffi / HLOOpt / cpu / PreRev 0.00003013 s 0.00001461602003473672 s 2.06
hlo_ffi / HLOOpt / cpu / PostRev 0.000029737 s 0.000016265079939330462 s 1.83
hlo_ffi / HLOOpt / cpu / BothRev 0.000029687 s 0.000014199220004229574 s 2.09
hlo_ffi / PartOpt / cpu / PreRev 0.000029875 s 0.000014724999928148464 s 2.03
hlo_ffi / PartOpt / cpu / PostRev 0.000029881 s 0.00001431689992386964 s 2.09
hlo_ffi / PartOpt / cpu / BothRev 0.000029742 s 0.000013945200025773374 s 2.13
hlo_ffi / IPartOpt / cpu / PreRev 0.000029454 s 0.000014579059989046071 s 2.02
hlo_ffi / IPartOpt / cpu / PostRev 0.000029643 s 0.000014473160026682308 s 2.05
hlo_ffi / IPartOpt / cpu / BothRev 0.000028938 s 0.000013991260002512718 s 2.07
hlo_ffi / DefOpt / cpu / PreRev 0.000029869 s 0.00001468752001528628 s 2.03
hlo_ffi / DefOpt / cpu / PostRev 0.000029774 s 0.000014148140035104009 s 2.10
hlo_ffi / DefOpt / cpu / BothRev 0.00002985 s 0.00001415964003172121 s 2.11
hlo_ffi / IDefOpt / cpu / PreRev 0.000029653 s 0.000014535159989463864 s 2.04
hlo_ffi / IDefOpt / cpu / PostRev 0.000029687 s 0.000014654960068583024 s 2.03
hlo_ffi / IDefOpt / cpu / BothRev 0.00002983 s 0.000014252880046115024 s 2.09
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.0009204894000049 s 0.0009126427998126 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 / Jax / cpu / Primal 0.0009299280000277 s 0.0008927976001359 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 / Primal 0.0009880593999469 s 0.0010200202001215 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.0009399049999956 s 0.0009099765997234 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 / IPartOpt / cpu / Primal 0.0009516319999647 s 0.0009122600000409 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 / DefOpt / cpu / Primal 0.0009935719999702 s 0.0009517657999822 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / Primal 0.0009884951999765 s 0.0009415242000613 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 / cpu / Forward 0.0023024988000088 s 0.0022322189999613 s 1.03
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Forward 0.0025613666000253 s 0.0022640238001258 s 1.13
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Forward 0.0023166976000311 s 0.0022092698000051 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 / PartOpt / cpu / Forward 0.0022911809999641 s 0.0022217054000066 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 / IPartOpt / cpu / Forward 0.0024424687999953 s 0.0021557281999776 s 1.13
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cpu / Forward 0.0023895834000086 s 0.002132012400034 s 1.12
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.0023649107999972 s 0.0021401981999588 s 1.10
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.0053143104000128 s 0.0054930844002228 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 / JaXPipe / cpu / PostRev 0.005989218600007 s 0.0055573514002389 s 1.08
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cpu / BothRev 0.0055158616000198 s 0.0056951962000312 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 / Jax / cpu / BothRev 0.0055554185999653 s 0.005988877800155 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 / PreRev 0.0052290580000317 s 0.005326209800296 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / PostRev 0.0052442613999801 s 0.0057709669999894 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 / HLOOpt / cpu / BothRev 0.0056484529999806 s 0.004952465000133 s 1.14
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.0058299041999816 s 0.0056679671999518 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 / cpu / PostRev 0.0052686790000052 s 0.0052149561999613 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / BothRev 0.00561458760003 s 0.005356651799957 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 / cpu / PreRev 0.0056060635999983 s 0.005281492999893 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 / cpu / PostRev 0.0060029984000038 s 0.005752405199928 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 / IPartOpt / cpu / BothRev 0.0060786337999616 s 0.0042657144002077 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 / DefOpt / cpu / PreRev 0.005925623600001 s 0.0053249480000886 s 1.11
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.0057202738000341 s 0.0053774022000652 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 / cpu / BothRev 0.0062362961999951 s 0.0053912424000372 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 / IDefOpt / cpu / PreRev 0.0043314718000146 s 0.0052342957998916 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 / IDefOpt / cpu / PostRev 0.0055874297999935 s 0.0053787539998666 s 1.04
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / BothRev 0.0055301765999956 s 0.0056076900000334 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / Primal 0.000273535 s 0.000294782 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 / Jax / cuda / Primal 0.000273086 s 0.000295519 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 / HLOOpt / cuda / Primal 0.00028675 s 0.000300735 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 / PartOpt / cuda / Primal 0.000271775 s 0.000294431 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 / cuda / Primal 0.000272543 s 0.000294719 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 / cuda / Primal 0.000287423 s 0.000302302 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 / IDefOpt / cuda / Primal 0.0002866549999999 s 0.000300958 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 / cuda / Forward 0.000558045 s 0.000582397 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / Forward 0.000538334 s 0.000565949 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 / HLOOpt / cuda / Forward 0.000558013 s 0.000582525 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / Forward 0.000557373 s 0.0005827489999999 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / Forward 0.000557374 s 0.000581821 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / Forward 0.000555838 s 0.000582045 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 / IDefOpt / cuda / Forward 0.000557118 s 0.000582365 s 0.96
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / PreRev 0.001028219 s 0.001053018 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / PostRev 0.000987356 s 0.001010011 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / cuda / BothRev 0.001026171 s 0.001049755 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cuda / BothRev 0.00098854 s 0.001002806 s 0.99
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / PreRev 0.001013147 s 0.001033755 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / PostRev 0.001039355 s 0.001060987 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cuda / BothRev 0.001012764 s 0.001036155 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / PreRev 0.001027579 s 0.001047547 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / PostRev 0.000979004 s 0.000998715 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cuda / BothRev 0.001028635 s 0.0010492729999999 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / PreRev 0.001027227 s 0.001049147 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / PostRev 0.000977179 s 0.000997211 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cuda / BothRev 0.001026651 s 0.001048443 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / PreRev 0.0010230349999999 s 0.001050459 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 / DefOpt / cuda / PostRev 0.0009624589999999 s 0.000983675 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / DefOpt / cuda / BothRev 0.001024828 s 0.001050168 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / PreRev 0.001023867 s 0.00104982 s 0.98
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cuda / PostRev 0.001021787 s 0.001053819 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 / IDefOpt / cuda / BothRev 0.001023323 s 0.0010519309999999 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 / JaXPipe / tpu / Primal 0.00012388225 s 0.000124404 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.00012679825 s 0.000126576 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.00015269 s 0.00015267225 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.0001348332499999 s 0.00013367925 s 1.01
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / Primal 0.00013108275 s 0.0001313635 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.0001477719999999 s 0.0001480195 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.00015113825 s 0.0001510995 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.0002120615 s 0.00021246825 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.000260808 s 0.0002609435 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.00021225425 s 0.000212679 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.0002184465 s 0.0002185557499999 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.00021224025 s 0.00021224075 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.0002183985 s 0.00021877125 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.0002120185 s 0.0002124265 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.00035487875 s 0.000354156 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.0002562725 s 0.0002565564999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / JaXPipe / tpu / BothRev 0.00035528375 s 0.000353798 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.0002572775 s 0.00025646925 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / tpu / PreRev 0.000355127 s 0.0003540605 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.000292037 s 0.0002909392499999 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.00035504975 s 0.0003538199999999 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.00035631425 s 0.00035552825 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.00027163525 s 0.0002709835 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.000356355 s 0.00035510475 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.0003553565 s 0.00035378575 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.00027256425 s 0.0002722284999999 s 1.00
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / tpu / BothRev 0.0003549022499999 s 0.00035379125 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.0003586245 s 0.0003576025 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.0002832335 s 0.00028285975 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.00035854825 s 0.000357185 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.000357924 s 0.00035625625 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.00030191075 s 0.0003005229999999 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.0003575654999999 s 0.00035597025 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.002404688 s 0.0009126427998126 s 2.63
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / Primal 0.002342788 s 0.0008927976001359 s 2.62
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / HLOOpt / cpu / Primal 0.002581845 s 0.0010200202001215 s 2.53
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.002300927 s 0.0009099765997234 s 2.53
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.002448134 s 0.0009122600000409 s 2.68
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.0023392599999999 s 0.0009517657999822 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 / Primal 0.002735267 s 0.0009415242000613 s 2.91
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.006332988 s 0.0022322189999613 s 2.84
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.006333829 s 0.0022640238001258 s 2.80
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.006008807 s 0.0022092698000051 s 2.72
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.006186834 s 0.0022217054000066 s 2.78
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.006340809 s 0.0021557281999776 s 2.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 / Forward 0.006251461 s 0.002132012400034 s 2.93
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.005831304 s 0.0021401981999588 s 2.72
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.009342601 s 0.0054930844002228 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 / PostRev 0.009735776 s 0.0055573514002389 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 / JaXPipe / cpu / BothRev 0.009265215 s 0.0056951962000312 s 1.63
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / Jax / cpu / BothRev 0.0097168469999999 s 0.005988877800155 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 / HLOOpt / cpu / PreRev 0.009999414 s 0.005326209800296 s 1.88
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.008744001 s 0.0057709669999894 s 1.52
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.010252755 s 0.004952465000133 s 2.07
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / PreRev 0.009893981 s 0.0056679671999518 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 / PartOpt / cpu / PostRev 0.009353087 s 0.0052149561999613 s 1.79
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / PartOpt / cpu / BothRev 0.008641574 s 0.005356651799957 s 1.61
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PreRev 0.008615227 s 0.005281492999893 s 1.63
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IPartOpt / cpu / PostRev 0.009336483 s 0.005752405199928 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 / BothRev 0.00945091 s 0.0042657144002077 s 2.22
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.00994018 s 0.0053249480000886 s 1.87
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.008068101 s 0.0053774022000652 s 1.50
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.008692885 s 0.0053912424000372 s 1.61
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PreRev 0.00975957 s 0.0052342957998916 s 1.86
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / PostRev 0.0096697349999999 s 0.0053787539998666 s 1.80
llama_dim_288_hidden_dim_768_n_layers_6_n_heads_6_n_kv_heads_6_vocab_size_32000_seq_len_256 / IDefOpt / cpu / BothRev 0.009376899 s 0.0056076900000334 s 1.67
scatter_sum / JaXPipe / cpu / Primal 0.000007648139996945247 s 0.000008034660004341276 s 0.95
scatter_sum / Jax / cpu / Primal 0.000007728699990821042 s 0.000007401199927699053 s 1.04
scatter_sum / HLOOpt / cpu / Primal 0.000007885000009082433 s 0.00000782368006184697 s 1.01
scatter_sum / PartOpt / cpu / Primal 0.000007602200005294435 s 0.000007279739966179477 s 1.04
scatter_sum / IPartOpt / cpu / Primal 0.000007901999995283404 s 0.000008091340077953646 s 0.98
scatter_sum / DefOpt / cpu / Primal 0.000007577919993764226 s 0.000007440540011884877 s 1.02
scatter_sum / IDefOpt / cpu / Primal 0.000007438840009399428 s 0.000007552500064775813 s 0.98
scatter_sum / JaXPipe / cpu / Forward 0.000011841099994853722 s 0.000011487500087241642 s 1.03
scatter_sum / Jax / cpu / Forward 0.00001169769999478376 s 0.0000111186600952351 s 1.05
scatter_sum / HLOOpt / cpu / Forward 0.000012346199998773954 s 0.000011329380049573956 s 1.09
scatter_sum / PartOpt / cpu / Forward 0.000011232440001549547 s 0.00001174510010969243 s 0.96
scatter_sum / IPartOpt / cpu / Forward 0.000011790739999923972 s 0.00001193121990581858 s 0.99
scatter_sum / DefOpt / cpu / Forward 0.000012004119998891836 s 0.000011192940000910312 s 1.07
scatter_sum / IDefOpt / cpu / Forward 0.000011696939995999856 s 0.00001139780006269575 s 1.03
scatter_sum / JaXPipe / cpu / PreRev 0.000011903800004802178 s 0.000011128400001325644 s 1.07
scatter_sum / JaXPipe / cpu / PostRev 0.000011963380018187308 s 0.000011160660160385304 s 1.07
scatter_sum / JaXPipe / cpu / BothRev 0.000012373900003694871 s 0.000011471000107121654 s 1.08
scatter_sum / Jax / cpu / BothRev 0.000011918180009615753 s 0.0000114425599531387 s 1.04
scatter_sum / HLOOpt / cpu / PreRev 0.000012078180002390582 s 0.000011476399977254916 s 1.05
scatter_sum / HLOOpt / cpu / PostRev 0.00001400890001150401 s 0.000013494479953806148 s 1.04
scatter_sum / HLOOpt / cpu / BothRev 0.000012350479998985976 s 0.000011364520014467417 s 1.09
scatter_sum / PartOpt / cpu / PreRev 0.0000118873600013103 s 0.000011737500026356429 s 1.01
scatter_sum / PartOpt / cpu / PostRev 0.000011857200004214974 s 0.000010909779975918354 s 1.09
scatter_sum / PartOpt / cpu / BothRev 0.000011894899994331354 s 0.00001189136008179048 s 1.00
scatter_sum / IPartOpt / cpu / PreRev 0.000011825740011772723 s 0.000011497700052132132 s 1.03
scatter_sum / IPartOpt / cpu / PostRev 0.00001186107999046726 s 0.000011243499975535088 s 1.05
scatter_sum / IPartOpt / cpu / BothRev 0.00001142133999564976 s 0.00001136295997639536 s 1.01
scatter_sum / DefOpt / cpu / PreRev 0.000011913079997611933 s 0.00001163702010671841 s 1.02
scatter_sum / DefOpt / cpu / PostRev 0.000011770920004892104 s 0.000010856139979296132 s 1.08
scatter_sum / DefOpt / cpu / BothRev 0.00001186358001177723 s 0.0000110379801117233 s 1.07
scatter_sum / IDefOpt / cpu / PreRev 0.00001150291999692854 s 0.000010984399905282773 s 1.05
scatter_sum / IDefOpt / cpu / PostRev 0.000012066939993928828 s 0.000011226860078750178 s 1.07
scatter_sum / IDefOpt / cpu / BothRev 0.00001179714000045351 s 0.000011499479951453397 s 1.03
scatter_sum / JaXPipe / cuda / Primal 0.000014272 s 0.00001056 s 1.35
scatter_sum / Jax / cuda / Primal 0.00000992 s 0.000010656 s 0.93
scatter_sum / HLOOpt / cuda / Primal 0.000010176 s 0.000010624 s 0.96
scatter_sum / PartOpt / cuda / Primal 0.000011168 s 0.000010496 s 1.06
scatter_sum / IPartOpt / cuda / Primal 0.000010112 s 0.000010816 s 0.93
scatter_sum / DefOpt / cuda / Primal 0.000010048 s 0.000010816 s 0.93
scatter_sum / IDefOpt / cuda / Primal 0.000010207 s 0.000010432 s 0.98
scatter_sum / JaXPipe / cuda / Forward 0.0000184 s 0.000017247 s 1.07
scatter_sum / Jax / cuda / Forward 0.000017343 s 0.000017216 s 1.01
scatter_sum / HLOOpt / cuda / Forward 0.000016736 s 0.000017536 s 0.95
scatter_sum / PartOpt / cuda / Forward 0.000017632 s 0.000017375999999999998 s 1.01
scatter_sum / IPartOpt / cuda / Forward 0.000017312 s 0.000017696 s 0.98
scatter_sum / DefOpt / cuda / Forward 0.000017503 s 0.000017536 s 1.00
scatter_sum / IDefOpt / cuda / Forward 0.000017024 s 0.000017247999999999998 s 0.99
scatter_sum / JaXPipe / cuda / PreRev 0.000017216 s 0.000017472 s 0.99
scatter_sum / JaXPipe / cuda / PostRev 0.000017536 s 0.000017408 s 1.01
scatter_sum / JaXPipe / cuda / BothRev 0.000017247999999999998 s 0.00001744 s 0.99
scatter_sum / Jax / cuda / BothRev 0.000016927999999999998 s 0.000017696 s 0.96
scatter_sum / HLOOpt / cuda / PreRev 0.000016832 s 0.000017888000000000002 s 0.94
scatter_sum / HLOOpt / cuda / PostRev 0.000016768000000000003 s 0.00001712 s 0.98
scatter_sum / HLOOpt / cuda / BothRev 0.00001888 s 0.00001728 s 1.09
scatter_sum / PartOpt / cuda / PreRev 0.000019648 s 0.000017312 s 1.13
scatter_sum / PartOpt / cuda / PostRev 0.000018304 s 0.000016767 s 1.09
scatter_sum / PartOpt / cuda / BothRev 0.000019072 s 0.000017344 s 1.10
scatter_sum / IPartOpt / cuda / PreRev 0.00001904 s 0.00001744 s 1.09
scatter_sum / IPartOpt / cuda / PostRev 0.000016896000000000002 s 0.00001712 s 0.99
scatter_sum / IPartOpt / cuda / BothRev 0.00001696 s 0.000017503999999999997 s 0.97
scatter_sum / DefOpt / cuda / PreRev 0.000016576000000000002 s 0.000017791 s 0.93
scatter_sum / DefOpt / cuda / PostRev 0.000019168 s 0.000016832 s 1.14
scatter_sum / DefOpt / cuda / BothRev 0.000016864 s 0.000017375999999999998 s 0.97
scatter_sum / IDefOpt / cuda / PreRev 0.000017408 s 0.000017824 s 0.98
scatter_sum / IDefOpt / cuda / PostRev 0.00001696 s 0.000017729000000000003 s 0.96
scatter_sum / IDefOpt / cuda / BothRev 0.000017375999999999998 s 0.000017312 s 1.00
scatter_sum / JaXPipe / tpu / Primal 0.0000013443250000000005 s 0.00000137035 s 0.98
scatter_sum / Jax / tpu / Primal 0.0000013442 s 0.000001343475 s 1.00
scatter_sum / HLOOpt / tpu / Primal 0.00000134325 s 0.000001371175 s 0.98
scatter_sum / PartOpt / tpu / Primal 0.0000013438249999999998 s 0.00000134325 s 1.00
scatter_sum / IPartOpt / tpu / Primal 0.0000013441249999999998 s 0.0000013707749999999998 s 0.98
scatter_sum / DefOpt / tpu / Primal 0.00000134365 s 0.000001344475 s 1.00
scatter_sum / IDefOpt / tpu / Primal 0.00000134435 s 0.0000013708 s 0.98
scatter_sum / JaXPipe / tpu / Forward 0.0000027424 s 0.000002754575 s 1.00
scatter_sum / Jax / tpu / Forward 0.000002746925 s 0.00000275575 s 1.00
scatter_sum / HLOOpt / tpu / Forward 0.000002740375 s 0.000002756875 s 0.99
scatter_sum / PartOpt / tpu / Forward 0.000002713275 s 0.00000272645 s 1.00
scatter_sum / IPartOpt / tpu / Forward 0.00000273785 s 0.0000027535 s 0.99
scatter_sum / DefOpt / tpu / Forward 0.000002713225 s 0.000002719975 s 1.00
scatter_sum / IDefOpt / tpu / Forward 0.000002736575 s 0.0000027513 s 0.99
scatter_sum / JaXPipe / tpu / PreRev 0.000002710475 s 0.000002714775 s 1.00
scatter_sum / JaXPipe / tpu / PostRev 0.00000273335 s 0.0000027531250000000004 s 0.99
scatter_sum / JaXPipe / tpu / BothRev 0.0000027256 s 0.000002727975 s 1.00
scatter_sum / Jax / tpu / BothRev 0.0000027855000000000003 s 0.0000028086750000000003 s 0.99
scatter_sum / HLOOpt / tpu / PreRev 0.000002728725 s 0.00000274125 s 1.00
scatter_sum / HLOOpt / tpu / PostRev 0.0000027893250000000006 s 0.0000027999 s 1.00
scatter_sum / HLOOpt / tpu / BothRev 0.0000027300750000000003 s 0.000002731275 s 1.00
scatter_sum / PartOpt / tpu / PreRev 0.0000027897 s 0.000002804625 s 0.99
scatter_sum / PartOpt / tpu / PostRev 0.000002722 s 0.0000027357500000000004 s 0.99
scatter_sum / PartOpt / tpu / BothRev 0.000002790625 s 0.0000028022000000000005 s 1.00
scatter_sum / IPartOpt / tpu / PreRev 0.00000273055 s 0.000002732275 s 1.00
scatter_sum / IPartOpt / tpu / PostRev 0.000002786125 s 0.0000028057 s 0.99
scatter_sum / IPartOpt / tpu / BothRev 0.000002730875 s 0.0000027315 s 1.00
scatter_sum / DefOpt / tpu / PreRev 0.00000278995 s 0.0000028033249999999995 s 1.00
scatter_sum / DefOpt / tpu / PostRev 0.000002720675 s 0.0000027280250000000004 s 1.00
scatter_sum / DefOpt / tpu / BothRev 0.000002791275 s 0.00000280605 s 0.99
scatter_sum / IDefOpt / tpu / PreRev 0.0000027266 s 0.00000273425 s 1.00
scatter_sum / IDefOpt / tpu / PostRev 0.000002786225 s 0.0000028042 s 0.99
scatter_sum / IDefOpt / tpu / BothRev 0.0000027299250000000004 s 0.0000027361 s 1.00
scatter_sum / JaXPipe / cpu / Primal 0.000019165 s 0.000008034660004341276 s 2.39
scatter_sum / Jax / cpu / Primal 0.000019336 s 0.000007401199927699053 s 2.61
scatter_sum / HLOOpt / cpu / Primal 0.00001937 s 0.00000782368006184697 s 2.48
scatter_sum / PartOpt / cpu / Primal 0.000018874 s 0.000007279739966179477 s 2.59
scatter_sum / IPartOpt / cpu / Primal 0.000018748 s 0.000008091340077953646 s 2.32
scatter_sum / DefOpt / cpu / Primal 0.000018974 s 0.000007440540011884877 s 2.55
scatter_sum / IDefOpt / cpu / Primal 0.000019457 s 0.000007552500064775813 s 2.58
scatter_sum / JaXPipe / cpu / Forward 0.000027552 s 0.000011487500087241642 s 2.40
scatter_sum / Jax / cpu / Forward 0.000027612000000000003 s 0.0000111186600952351 s 2.48
scatter_sum / HLOOpt / cpu / Forward 0.000027309 s 0.000011329380049573956 s 2.41
scatter_sum / PartOpt / cpu / Forward 0.000027435 s 0.00001174510010969243 s 2.34
scatter_sum / IPartOpt / cpu / Forward 0.000028068 s 0.00001193121990581858 s 2.35
scatter_sum / DefOpt / cpu / Forward 0.000034221 s 0.000011192940000910312 s 3.06
scatter_sum / IDefOpt / cpu / Forward 0.000027201 s 0.00001139780006269575 s 2.39
scatter_sum / JaXPipe / cpu / PreRev 0.000027776 s 0.000011128400001325644 s 2.50
scatter_sum / JaXPipe / cpu / PostRev 0.000027909 s 0.000011160660160385304 s 2.50
scatter_sum / JaXPipe / cpu / BothRev 0.000028105000000000003 s 0.000011471000107121654 s 2.45
scatter_sum / Jax / cpu / BothRev 0.000027957 s 0.0000114425599531387 s 2.44
scatter_sum / HLOOpt / cpu / PreRev 0.000027671 s 0.000011476399977254916 s 2.41
scatter_sum / HLOOpt / cpu / PostRev 0.000028404 s 0.000013494479953806148 s 2.10
scatter_sum / HLOOpt / cpu / BothRev 0.000033508000000000005 s 0.000011364520014467417 s 2.95
scatter_sum / PartOpt / cpu / PreRev 0.000027697 s 0.000011737500026356429 s 2.36
scatter_sum / PartOpt / cpu / PostRev 0.000033713 s 0.000010909779975918354 s 3.09
scatter_sum / PartOpt / cpu / BothRev 0.000027388 s 0.00001189136008179048 s 2.30
scatter_sum / IPartOpt / cpu / PreRev 0.00002808 s 0.000011497700052132132 s 2.44
scatter_sum / IPartOpt / cpu / PostRev 0.000027102 s 0.000011243499975535088 s 2.41
scatter_sum / IPartOpt / cpu / BothRev 0.000027757 s 0.00001136295997639536 s 2.44
scatter_sum / DefOpt / cpu / PreRev 0.00002782 s 0.00001163702010671841 s 2.39
scatter_sum / DefOpt / cpu / PostRev 0.000027837 s 0.000010856139979296132 s 2.56
scatter_sum / DefOpt / cpu / BothRev 0.000028384 s 0.0000110379801117233 s 2.57
scatter_sum / IDefOpt / cpu / PreRev 0.000027918 s 0.000010984399905282773 s 2.54
scatter_sum / IDefOpt / cpu / PostRev 0.000027558 s 0.000011226860078750178 s 2.45
scatter_sum / IDefOpt / cpu / BothRev 0.000027801 s 0.000011499479951453397 s 2.42
slicing / JaXPipe / cpu / Primal 0.000006109079997713706 s 0.000006256359974941006 s 0.98
slicing / Jax / cpu / Primal 0.000006085980000989366 s 0.000006476919879787601 s 0.94
slicing / HLOOpt / cpu / Primal 0.000006277560000853555 s 0.000006179880001582205 s 1.02
slicing / PartOpt / cpu / Primal 0.000006613360014853243 s 0.000006054379955457989 s 1.09
slicing / IPartOpt / cpu / Primal 0.000007016240006123553 s 0.000006397199995262781 s 1.10
slicing / DefOpt / cpu / Primal 0.000006453379996855801 s 0.000005949320056970464 s 1.08
slicing / IDefOpt / cpu / Primal 0.0000065935599968725 s 0.000006439940079872031 s 1.02
slicing / JaXPipe / cpu / Forward 0.000009579260006375989 s 0.000009397520007041748 s 1.02
slicing / Jax / cpu / Forward 0.000010502780003207593 s 0.000009175039958790875 s 1.14
slicing / HLOOpt / cpu / Forward 0.000009791460001906673 s 0.00000989017997198971 s 0.99
slicing / PartOpt / cpu / Forward 0.000010003699990193128 s 0.00000902590001714998 s 1.11
slicing / IPartOpt / cpu / Forward 0.000009670279987403771 s 0.000009671999978309033 s 1.00
slicing / DefOpt / cpu / Forward 0.0000097823999863067 s 0.00000922690007428173 s 1.06
slicing / IDefOpt / cpu / Forward 0.000009925180006575828 s 0.000009310740024375264 s 1.07
slicing / JaXPipe / cpu / PreRev 0.000010111639994647705 s 0.00001027858004817972 s 0.98
slicing / JaXPipe / cpu / PostRev 0.000010164660013742832 s 0.000009631540033296916 s 1.06
slicing / JaXPipe / cpu / BothRev 0.00001073605998954008 s 0.000010070120042655615 s 1.07
slicing / Jax / cpu / BothRev 0.000010315040003661123 s 0.00000994826003079652 s 1.04
slicing / HLOOpt / cpu / PreRev 0.000010879500002829446 s 0.000009710159956739516 s 1.12
slicing / HLOOpt / cpu / PostRev 0.000011923799997930472 s 0.000011341420031385496 s 1.05
slicing / HLOOpt / cpu / BothRev 0.000009621160006645367 s 0.00000971812001807848 s 0.99
slicing / PartOpt / cpu / PreRev 0.00001016325999671608 s 0.00001009372001135489 s 1.01
slicing / PartOpt / cpu / PostRev 0.00000981477999175695 s 0.000009648800096329067 s 1.02
slicing / PartOpt / cpu / BothRev 0.000010438339995744172 s 0.000010046099978353596 s 1.04
slicing / IPartOpt / cpu / PreRev 0.000010103979996074486 s 0.000009653820034145613 s 1.05
slicing / IPartOpt / cpu / PostRev 0.00001045937999151647 s 0.000010048899948742472 s 1.04
slicing / IPartOpt / cpu / BothRev 0.000009766619996298688 s 0.000010373479999543634 s 0.94
slicing / DefOpt / cpu / PreRev 0.000010189059996719153 s 0.000009390980030730134 s 1.08
slicing / DefOpt / cpu / PostRev 0.000010399799996321236 s 0.000009702339957584628 s 1.07
slicing / DefOpt / cpu / BothRev 0.00001031676001048254 s 0.00000979103999270592 s 1.05
slicing / IDefOpt / cpu / PreRev 0.000009986280008433824 s 0.000009572320104780374 s 1.04
slicing / IDefOpt / cpu / PostRev 0.000009948079991772827 s 0.000010150820035050856 s 0.98
slicing / IDefOpt / cpu / BothRev 0.00001010718000770794 s 0.000009666499990999 s 1.05
slicing / JaXPipe / cuda / Primal 0.000001887 s 0.000002303 s 0.82
slicing / Jax / cuda / Primal 0.000001887 s 0.000002304 s 0.82
slicing / HLOOpt / cuda / Primal 0.000001887 s 0.000002303 s 0.82
slicing / PartOpt / cuda / Primal 0.000001887 s 0.000002303 s 0.82
slicing / IPartOpt / cuda / Primal 0.000001887 s 0.000002303 s 0.82
slicing / DefOpt / cuda / Primal 0.000001887 s 0.000002303 s 0.82
slicing / IDefOpt / cuda / Primal 0.000001888 s 0.000002304 s 0.82
slicing / JaXPipe / cuda / Forward 0.00001024 s 0.000010272 s 1.00
slicing / Jax / cuda / Forward 0.000010272 s 0.00001008 s 1.02
slicing / HLOOpt / cuda / Forward 0.000010144 s 0.000010304 s 0.98
slicing / PartOpt / cuda / Forward 0.000010016 s 0.000010592 s 0.95
slicing / IPartOpt / cuda / Forward 0.000010144 s 0.00000992 s 1.02
slicing / DefOpt / cuda / Forward 0.000010176 s 0.00001024 s 0.99
slicing / IDefOpt / cuda / Forward 0.000009888 s 0.00001088 s 0.91
slicing / JaXPipe / cuda / PreRev 0.000010048 s 0.0000104 s 0.97
slicing / JaXPipe / cuda / PostRev 0.000009888 s 0.000010336 s 0.96
slicing / JaXPipe / cuda / BothRev 0.000010048 s 0.0000104 s 0.97
slicing / Jax / cuda / BothRev 0.000010272 s 0.000010208 s 1.01
slicing / HLOOpt / cuda / PreRev 0.000010175 s 0.000010528 s 0.97
slicing / HLOOpt / cuda / PostRev 0.000010049 s 0.000010591 s 0.95
slicing / HLOOpt / cuda / BothRev 0.000010016 s 0.000010496 s 0.95
slicing / PartOpt / cuda / PreRev 0.000009664 s 0.000010368 s 0.93
slicing / PartOpt / cuda / PostRev 0.000010304 s 0.0000104 s 0.99
slicing / PartOpt / cuda / BothRev 0.000010336 s 0.000010815 s 0.96
slicing / IPartOpt / cuda / PreRev 0.000010048 s 0.000010912 s 0.92
slicing / IPartOpt / cuda / PostRev 0.000009568 s 0.00001072 s 0.89
slicing / IPartOpt / cuda / BothRev 0.000009984 s 0.000010368 s 0.96
slicing / DefOpt / cuda / PreRev 0.000010176 s 0.00001024 s 0.99
slicing / DefOpt / cuda / PostRev 0.000009952 s 0.000010432 s 0.95
slicing / DefOpt / cuda / BothRev 0.000010144 s 0.000010496 s 0.97
slicing / IDefOpt / cuda / PreRev 0.000010208 s 0.000010464 s 0.98
slicing / IDefOpt / cuda / PostRev 0.000009985 s 0.000010656 s 0.94
slicing / IDefOpt / cuda / BothRev 0.000009984 s 0.000010305 s 0.97
slicing / JaXPipe / tpu / Primal 9.50125e-7 s 0.00000103045 s 0.92
slicing / Jax / tpu / Primal 9.541e-7 s 9.65925e-7 s 0.99
slicing / HLOOpt / tpu / Primal 9.51525e-7 s 0.000001024375 s 0.93
slicing / PartOpt / tpu / Primal 9.51575e-7 s 9.59525e-7 s 0.99
slicing / IPartOpt / tpu / Primal 9.5345e-7 s 0.0000010249250000000002 s 0.93
slicing / DefOpt / tpu / Primal 9.51475e-7 s 9.62125e-7 s 0.99
slicing / IDefOpt / tpu / Primal 9.558e-7 s 0.000001022275 s 0.93
slicing / JaXPipe / tpu / Forward 0.0000014038999999999998 s 0.000001408625 s 1.00
slicing / Jax / tpu / Forward 0.0000014051749999999998 s 0.000001476525 s 0.95
slicing / HLOOpt / tpu / Forward 0.00000151685 s 0.0000015182499999999998 s 1.00
slicing / PartOpt / tpu / Forward 0.000001419575 s 0.0000014921 s 0.95
slicing / IPartOpt / tpu / Forward 0.000001510375 s 0.000001516925 s 1.00
slicing / DefOpt / tpu / Forward 0.000001425 s 0.00000149725 s 0.95
slicing / IDefOpt / tpu / Forward 0.000001511825 s 0.0000015184500000000005 s 1.00
slicing / JaXPipe / tpu / PreRev 0.00000233835 s 0.0000025648 s 0.91
slicing / JaXPipe / tpu / PostRev 0.000002510725 s 0.000002508975 s 1.00
slicing / JaXPipe / tpu / BothRev 0.00000235895 s 0.0000025731 s 0.92
slicing / Jax / tpu / BothRev 0.0000025238000000000003 s 0.00000252785 s 1.00
slicing / HLOOpt / tpu / PreRev 0.000002343675 s 0.00000257755 s 0.91
slicing / HLOOpt / tpu / PostRev 0.000002537675 s 0.00000252875 s 1.00
slicing / HLOOpt / tpu / BothRev 0.0000023551 s 0.0000025746000000000003 s 0.91
slicing / PartOpt / tpu / PreRev 0.000002517475 s 0.000002530725 s 0.99
slicing / PartOpt / tpu / PostRev 0.0000023624 s 0.0000025674 s 0.92
slicing / PartOpt / tpu / BothRev 0.00000252845 s 0.000002532 s 1.00
slicing / IPartOpt / tpu / PreRev 0.000002344175 s 0.000002570325 s 0.91
slicing / IPartOpt / tpu / PostRev 0.00000254495 s 0.000002529525 s 1.01
slicing / IPartOpt / tpu / BothRev 0.0000023471 s 0.0000025746000000000003 s 0.91
slicing / DefOpt / tpu / PreRev 0.000002529025 s 0.000002531425 s 1.00
slicing / DefOpt / tpu / PostRev 0.0000023538 s 0.000002575075 s 0.91
slicing / DefOpt / tpu / BothRev 0.0000025178000000000003 s 0.000002524525 s 1.00
slicing / IDefOpt / tpu / PreRev 0.00000235235 s 0.000002573175 s 0.91
slicing / IDefOpt / tpu / PostRev 0.000002528525 s 0.000002527475 s 1.00
slicing / IDefOpt / tpu / BothRev 0.00000236015 s 0.000002570325 s 0.92
slicing / JaXPipe / cpu / Primal 0.000015784 s 0.000006256359974941006 s 2.52
slicing / Jax / cpu / Primal 0.000015502 s 0.000006476919879787601 s 2.39
slicing / HLOOpt / cpu / Primal 0.000015644 s 0.000006179880001582205 s 2.53
slicing / PartOpt / cpu / Primal 0.000015336 s 0.000006054379955457989 s 2.53
slicing / IPartOpt / cpu / Primal 0.000015779999999999998 s 0.000006397199995262781 s 2.47
slicing / DefOpt / cpu / Primal 0.000015476 s 0.000005949320056970464 s 2.60
slicing / IDefOpt / cpu / Primal 0.000015408 s 0.000006439940079872031 s 2.39
slicing / JaXPipe / cpu / Forward 0.000021341 s 0.000009397520007041748 s 2.27
slicing / Jax / cpu / Forward 0.000020788 s 0.000009175039958790875 s 2.27
slicing / HLOOpt / cpu / Forward 0.000020685 s 0.00000989017997198971 s 2.09
slicing / PartOpt / cpu / Forward 0.000020822 s 0.00000902590001714998 s 2.31
slicing / IPartOpt / cpu / Forward 0.000026548 s 0.000009671999978309033 s 2.74
slicing / DefOpt / cpu / Forward 0.000020965 s 0.00000922690007428173 s 2.27
slicing / IDefOpt / cpu / Forward 0.000021107 s 0.000009310740024375264 s 2.27
slicing / JaXPipe / cpu / PreRev 0.000021626 s 0.00001027858004817972 s 2.10
slicing / JaXPipe / cpu / PostRev 0.000021894 s 0.000009631540033296916 s 2.27
slicing / JaXPipe / cpu / BothRev 0.000021637 s 0.000010070120042655615 s 2.15
slicing / Jax / cpu / BothRev 0.000021869 s 0.00000994826003079652 s 2.20
slicing / HLOOpt / cpu / PreRev 0.000021476 s 0.000009710159956739516 s 2.21
slicing / HLOOpt / cpu / PostRev 0.000022063 s 0.000011341420031385496 s 1.95
slicing / HLOOpt / cpu / BothRev 0.000027392 s 0.00000971812001807848 s 2.82
slicing / PartOpt / cpu / PreRev 0.000021422 s 0.00001009372001135489 s 2.12
slicing / PartOpt / cpu / PostRev 0.000021622 s 0.000009648800096329067 s 2.24
slicing / PartOpt / cpu / BothRev 0.000021825 s 0.000010046099978353596 s 2.17
slicing / IPartOpt / cpu / PreRev 0.000021468 s 0.000009653820034145613 s 2.22
slicing / IPartOpt / cpu / PostRev 0.000021498 s 0.000010048899948742472 s 2.14
slicing / IPartOpt / cpu / BothRev 0.000022181 s 0.000010373479999543634 s 2.14
slicing / DefOpt / cpu / PreRev 0.000021374 s 0.000009390980030730134 s 2.28
slicing / DefOpt / cpu / PostRev 0.000021342 s 0.000009702339957584628 s 2.20
slicing / DefOpt / cpu / BothRev 0.000021532 s 0.00000979103999270592 s 2.20
slicing / IDefOpt / cpu / PreRev 0.000021593 s 0.000009572320104780374 s 2.26
slicing / IDefOpt / cpu / PostRev 0.000021735 s 0.000010150820035050856 s 2.14
slicing / IDefOpt / cpu / BothRev 0.00002128 s 0.000009666499990999 s 2.20
sum / JaXPipe / cpu / Primal 0.000008197720001135167 s 0.000007971080012794119 s 1.03
sum / Jax / cpu / Primal 0.00000781075998020242 s 0.000007625819998793304 s 1.02
sum / HLOOpt / cpu / Primal 0.000008021799994821776 s 0.000008017239943001186 s 1.00
sum / PartOpt / cpu / Primal 0.00000754313999095757 s 0.000007410100042761769 s 1.02
sum / IPartOpt / cpu / Primal 0.000007949360003749462 s 0.000008181239918485517 s 0.97
sum / DefOpt / cpu / Primal 0.000007664219995149323 s 0.000007373279931925935 s 1.04
sum / IDefOpt / cpu / Primal 0.000007506980005018704 s 0.000007935899957374203 s 0.95
sum / JaXPipe / cpu / Forward 0.000011410420006541244 s 0.000011452859980636275 s 1.00
sum / Jax / cpu / Forward 0.00001165782001635307 s 0.0000106855801641359 s 1.09
sum / HLOOpt / cpu / Forward 0.000011438939995969122 s 0.000011622879992501113 s 0.98
sum / PartOpt / cpu / Forward 0.000011433360002683912 s 0.000011153219929838087 s 1.03
sum / IPartOpt / cpu / Forward 0.000011695699993197194 s 0.00001095907997296308 s 1.07
sum / DefOpt / cpu / Forward 0.000011900059992058231 s 0.000011301559970888776 s 1.05
sum / IDefOpt / cpu / Forward 0.000011227440002130606 s 0.000011096039925178047 s 1.01
sum / JaXPipe / cpu / PreRev 0.000011254579994783852 s 0.00001128660000176751 s 1.00
sum / JaXPipe / cpu / PostRev 0.000011664379997000651 s 0.000010986259912897368 s 1.06
sum / JaXPipe / cpu / BothRev 0.000011874579997765978 s 0.000011202440018678316 s 1.06
sum / Jax / cpu / BothRev 0.000011237159992560918 s 0.000010707339952205074 s 1.05
sum / HLOOpt / cpu / PreRev 0.000011166679996676976 s 0.000011084920060966397 s 1.01
sum / HLOOpt / cpu / PostRev 0.000013245179991372424 s 0.000013082000023132423 s 1.01
sum / HLOOpt / cpu / BothRev 0.000010821580017363886 s 0.000010698379956011197 s 1.01
sum / PartOpt / cpu / PreRev 0.000010647100004916864 s 0.000010895840041484915 s 0.98
sum / PartOpt / cpu / PostRev 0.000011103179990641366 s 0.000010930300031759544 s 1.02
sum / PartOpt / cpu / BothRev 0.000011547279991646064 s 0.000011242800028412602 s 1.03
sum / IPartOpt / cpu / PreRev 0.0000112957200008168 s 0.000010910799974226392 s 1.04
sum / IPartOpt / cpu / PostRev 0.00001062042000285146 s 0.00001021207994199358 s 1.04
sum / IPartOpt / cpu / BothRev 0.000011321540007429576 s 0.000010417019948363305 s 1.09
sum / DefOpt / cpu / PreRev 0.000011102840005605684 s 0.00001097947999369353 s 1.01
sum / DefOpt / cpu / PostRev 0.000010816480003086326 s 0.000010882719961955444 s 0.99
sum / DefOpt / cpu / BothRev 0.00001114339999276126 s 0.00001088356009859126 s 1.02
sum / IDefOpt / cpu / PreRev 0.00001106619999063696 s 0.000011024900013580918 s 1.00
sum / IDefOpt / cpu / PostRev 0.000011343799999394832 s 0.000010991760118486127 s 1.03
sum / IDefOpt / cpu / BothRev 0.000010844580003777082 s 0.000010506819999136496 s 1.03
sum / JaXPipe / cuda / Primal 0.000002047 s 0.000002464 s 0.83
sum / Jax / cuda / Primal 0.000002048 s 0.000002463 s 0.83
sum / HLOOpt / cuda / Primal 0.000002047 s 0.000002463 s 0.83
sum / PartOpt / cuda / Primal 0.000002047 s 0.000002463 s 0.83
sum / IPartOpt / cuda / Primal 0.000002048 s 0.000002463 s 0.83
sum / DefOpt / cuda / Primal 0.000002048 s 0.000002463 s 0.83
sum / IDefOpt / cuda / Primal 0.000002048 s 0.000002463 s 0.83
sum / JaXPipe / cuda / Forward 0.000010688 s 0.000010623 s 1.01
sum / Jax / cuda / Forward 0.000009888 s 0.000010464 s 0.94
sum / HLOOpt / cuda / Forward 0.000010272 s 0.000009952 s 1.03
sum / PartOpt / cuda / Forward 0.000010432 s 0.000010592 s 0.98
sum / IPartOpt / cuda / Forward 0.000010592 s 0.000010496 s 1.01
sum / DefOpt / cuda / Forward 0.00001056 s 0.000010464 s 1.01
sum / IDefOpt / cuda / Forward 0.000010368 s 0.000010496 s 0.99
sum / JaXPipe / cuda / PreRev 0.000010464 s 0.00001024 s 1.02
sum / JaXPipe / cuda / PostRev 0.000009696 s 0.000010592 s 0.92
sum / JaXPipe / cuda / BothRev 0.000009536 s 0.000010304 s 0.93
sum / Jax / cuda / BothRev 0.00001024 s 0.000010208 s 1.00
sum / HLOOpt / cuda / PreRev 0.000009856 s 0.000010432 s 0.94
sum / HLOOpt / cuda / PostRev 0.000008959999999999999 s 0.000010368 s 0.86
sum / HLOOpt / cuda / BothRev 0.000009728 s 0.000010111 s 0.96
sum / PartOpt / cuda / PreRev 0.000009568 s 0.0000104 s 0.92
sum / PartOpt / cuda / PostRev 0.00001024 s 0.000010272 s 1.00
sum / PartOpt / cuda / BothRev 0.000015199 s 0.00000944 s 1.61
sum / IPartOpt / cuda / PreRev 0.000010016 s 0.000010496 s 0.95
sum / IPartOpt / cuda / PostRev 0.000009888 s 0.000010272 s 0.96
sum / IPartOpt / cuda / BothRev 0.000009984 s 0.000010368 s 0.96
sum / DefOpt / cuda / PreRev 0.000010112 s 0.00001088 s 0.93
sum / DefOpt / cuda / PostRev 0.000009888 s 0.000010976 s 0.90
sum / DefOpt / cuda / BothRev 0.000010848 s 0.000010592 s 1.02
sum / IDefOpt / cuda / PreRev 0.000010624 s 0.000010944 s 0.97
sum / IDefOpt / cuda / PostRev 0.000010528 s 0.000010336 s 1.02
sum / IDefOpt / cuda / BothRev 0.00000992 s 0.000010656 s 0.93
sum / JaXPipe / tpu / Primal 5.031e-7 s 5.03125e-7 s 1.00
sum / Jax / tpu / Primal 5.466e-7 s 5.4745e-7 s 1.00
sum / HLOOpt / tpu / Primal 5.034e-7 s 5.036e-7 s 1.00
sum / PartOpt / tpu / Primal 5.472250000000001e-7 s 5.47275e-7 s 1.00
sum / IPartOpt / tpu / Primal 5.02775e-7 s 5.0345e-7 s 1.00
sum / DefOpt / tpu / Primal 5.471999999999999e-7 s 5.47275e-7 s 1.00
sum / IDefOpt / tpu / Primal 5.02975e-7 s 5.0305e-7 s 1.00
sum / JaXPipe / tpu / Forward 0.00000155805 s 0.000001551925 s 1.00
sum / Jax / tpu / Forward 0.0000014975749999999998 s 0.00000149655 s 1.00
sum / HLOOpt / tpu / Forward 0.00000153565 s 0.000001529625 s 1.00
sum / PartOpt / tpu / Forward 0.000001491275 s 0.00000149105 s 1.00
sum / IPartOpt / tpu / Forward 0.00000153495 s 0.00000152995 s 1.00
sum / DefOpt / tpu / Forward 0.0000014881 s 0.0000014889749999999998 s 1.00
sum / IDefOpt / tpu / Forward 0.0000015307999999999998 s 0.0000015351 s 1.00
sum / JaXPipe / tpu / PreRev 9.87075e-7 s 0.0000010405 s 0.95
sum / JaXPipe / tpu / PostRev 0.000001030975 s 0.0000010893 s 0.95
sum / JaXPipe / tpu / BothRev 9.92425e-7 s 0.000001042775 s 0.95
sum / Jax / tpu / BothRev 0.000001039775 s 0.000001087625 s 0.96
sum / HLOOpt / tpu / PreRev 9.9365e-7 s 0.00000104035 s 0.96
sum / HLOOpt / tpu / PostRev 0.0000010335 s 0.0000010899499999999998 s 0.95
sum / HLOOpt / tpu / BothRev 9.91425e-7 s 0.00000103875 s 0.95
sum / PartOpt / tpu / PreRev 0.000001033825 s 0.0000010843 s 0.95
sum / PartOpt / tpu / PostRev 9.8915e-7 s 0.000001041575 s 0.95
sum / PartOpt / tpu / BothRev 0.00000103135 s 0.000001088925 s 0.95
sum / IPartOpt / tpu / PreRev 9.85375e-7 s 0.000001046425 s 0.94
sum / IPartOpt / tpu / PostRev 0.0000010397000000000002 s 0.0000010828999999999998 s 0.96
sum / IPartOpt / tpu / BothRev 9.86675e-7 s 0.000001040425 s 0.95
sum / DefOpt / tpu / PreRev 0.0000010325 s 0.000001083325 s 0.95
sum / DefOpt / tpu / PostRev 9.929749999999998e-7 s 0.000001040375 s 0.95
sum / DefOpt / tpu / BothRev 0.00000105265 s 0.000001094025 s 0.96
sum / IDefOpt / tpu / PreRev 9.874e-7 s 0.0000010473499999999998 s 0.94
sum / IDefOpt / tpu / PostRev 0.00000103705 s 0.0000010966750000000002 s 0.95
sum / IDefOpt / tpu / BothRev 9.966e-7 s 0.0000010399750000000002 s 0.96
sum / JaXPipe / cpu / Primal 0.000018417 s 0.000007971080012794119 s 2.31
sum / Jax / cpu / Primal 0.000017663 s 0.000007625819998793304 s 2.32
sum / HLOOpt / cpu / Primal 0.000017831 s 0.000008017239943001186 s 2.22
sum / PartOpt / cpu / Primal 0.000017642 s 0.000007410100042761769 s 2.38
sum / IPartOpt / cpu / Primal 0.000017833 s 0.000008181239918485517 s 2.18
sum / DefOpt / cpu / Primal 0.000017636 s 0.000007373279931925935 s 2.39
sum / IDefOpt / cpu / Primal 0.000018268 s 0.000007935899957374203 s 2.30
sum / JaXPipe / cpu / Forward 0.000024667000000000003 s 0.000011452859980636275 s 2.15
sum / Jax / cpu / Forward 0.000024564 s 0.0000106855801641359 s 2.30
sum / HLOOpt / cpu / Forward 0.000024297 s 0.000011622879992501113 s 2.09
sum / PartOpt / cpu / Forward 0.000024434 s 0.000011153219929838087 s 2.19
sum / IPartOpt / cpu / Forward 0.000025061 s 0.00001095907997296308 s 2.29
sum / DefOpt / cpu / Forward 0.000024532 s 0.000011301559970888776 s 2.17
sum / IDefOpt / cpu / Forward 0.000024557 s 0.000011096039925178047 s 2.21
sum / JaXPipe / cpu / PreRev 0.000024199 s 0.00001128660000176751 s 2.14
sum / JaXPipe / cpu / PostRev 0.000023247 s 0.000010986259912897368 s 2.12
sum / JaXPipe / cpu / BothRev 0.000023644 s 0.000011202440018678316 s 2.11
sum / Jax / cpu / BothRev 0.000023471 s 0.000010707339952205074 s 2.19
sum / HLOOpt / cpu / PreRev 0.000023565 s 0.000011084920060966397 s 2.13
sum / HLOOpt / cpu / PostRev 0.000023414 s 0.000013082000023132423 s 1.79
sum / HLOOpt / cpu / BothRev 0.000023524 s 0.000010698379956011197 s 2.20
sum / PartOpt / cpu / PreRev 0.000023422 s 0.000010895840041484915 s 2.15
sum / PartOpt / cpu / PostRev 0.000023482 s 0.000010930300031759544 s 2.15
sum / PartOpt / cpu / BothRev 0.000023568 s 0.000011242800028412602 s 2.10
sum / IPartOpt / cpu / PreRev 0.000023369 s 0.000010910799974226392 s 2.14
sum / IPartOpt / cpu / PostRev 0.000023249 s 0.00001021207994199358 s 2.28
sum / IPartOpt / cpu / BothRev 0.000023266 s 0.000010417019948363305 s 2.23
sum / DefOpt / cpu / PreRev 0.00002352 s 0.00001097947999369353 s 2.14
sum / DefOpt / cpu / PostRev 0.000023598 s 0.000010882719961955444 s 2.17
sum / DefOpt / cpu / BothRev 0.000023319 s 0.00001088356009859126 s 2.14
sum / IDefOpt / cpu / PreRev 0.000023092 s 0.000011024900013580918 s 2.09
sum / IDefOpt / cpu / PostRev 0.000023587 s 0.000010991760118486127 s 2.15
sum / IDefOpt / cpu / BothRev 0.000029046 s 0.000010506819999136496 s 2.76
value_and_grad / JaXPipe / cpu / Primal 0.00001401457999008926 s 0.000013755879972450204 s 1.02
value_and_grad / Jax / cpu / Primal 0.000014454100009970716 s 0.000013791200035484508 s 1.05
value_and_grad / HLOOpt / cpu / Primal 0.00001397760000827475 s 0.000013499440010491526 s 1.04
value_and_grad / PartOpt / cpu / Primal 0.000013165959994694277 s 0.00001282425997487735 s 1.03
value_and_grad / IPartOpt / cpu / Primal 0.000013100060002670945 s 0.0000132059399766149 s 0.99
value_and_grad / DefOpt / cpu / Primal 0.000014051720002044022 s 0.00001333509995674831 s 1.05
value_and_grad / IDefOpt / cpu / Primal 0.00001329064000174185 s 0.000013777679923805408 s 0.96
value_and_grad / JaXPipe / cuda / Primal 0.000033727 s 0.000033727 s 1
value_and_grad / Jax / cuda / Primal 0.000034303 s 0.000033375000000000005 s 1.03
value_and_grad / HLOOpt / cuda / Primal 0.000034432 s 0.000033632 s 1.02
value_and_grad / PartOpt / cuda / Primal 0.000033951 s 0.000033695 s 1.01
value_and_grad / IPartOpt / cuda / Primal 0.00003408 s 0.00003408 s 1
value_and_grad / DefOpt / cuda / Primal 0.0000432 s 0.000033759999999999995 s 1.28
value_and_grad / IDefOpt / cuda / Primal 0.000034111 s 0.000034208 s 1.00
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.000028639 s 0.000013755879972450204 s 2.08
value_and_grad / Jax / cpu / Primal 0.000027457 s 0.000013791200035484508 s 1.99
value_and_grad / HLOOpt / cpu / Primal 0.000027828 s 0.000013499440010491526 s 2.06
value_and_grad / PartOpt / cpu / Primal 0.000027959 s 0.00001282425997487735 s 2.18
value_and_grad / IPartOpt / cpu / Primal 0.000027685 s 0.0000132059399766149 s 2.10
value_and_grad / DefOpt / cpu / Primal 0.000027718 s 0.00001333509995674831 s 2.08
value_and_grad / IDefOpt / cpu / Primal 0.000028147 s 0.000013777679923805408 s 2.04
jaxmd20 / JaXPipe / cuda / Primal 0.001545561 s 0.001465816 s 1.05
jaxmd20 / Jax / cuda / Primal 0.001538648 s 0.001531639 s 1.00
jaxmd20 / HLOOpt / cuda / Primal 0.001302713 s 0.00137385 s 0.95
jaxmd20 / PartOpt / cuda / Primal 0.001379897 s 0.0013668089999999 s 1.01
jaxmd20 / IPartOpt / cuda / Primal 0.0013404419999999 s 0.00134812 s 0.99
jaxmd20 / DefOpt / cuda / Primal 0.00092246 s 0.000938491 s 0.98
jaxmd20 / IDefOpt / cuda / Primal 0.0009553239999999 s 0.000962875 s 0.99
jaxmd20 / JaXPipe / cuda / Forward 0.001561367 s 0.001631768 s 0.96
jaxmd20 / Jax / cuda / Forward 0.001820152 s 0.001852822 s 0.98
jaxmd20 / HLOOpt / cuda / Forward 0.001630136 s 0.001709686 s 0.95
jaxmd20 / PartOpt / cuda / Forward 0.001655415 s 0.001719518 s 0.96
jaxmd20 / IPartOpt / cuda / Forward 0.0016220079999999 s 0.0017157329999999 s 0.95
jaxmd20 / DefOpt / cuda / Forward 0.001644536 s 0.001707192 s 0.96
jaxmd20 / IDefOpt / cuda / Forward 0.001633943 s 0.0017267369999999 s 0.95
jaxmd20 / JaXPipe / cuda / PreRev 0.002683602 s 0.002786098 s 0.96
jaxmd20 / JaXPipe / cuda / PostRev 0.005363652 s 0.005518528 s 0.97
jaxmd20 / JaXPipe / cuda / BothRev 0.002693427 s 0.002824722 s 0.95
jaxmd20 / Jax / cuda / BothRev 0.0053409 s 0.005536183 s 0.96
jaxmd20 / HLOOpt / cuda / PreRev 0.002752628 s 0.002873778 s 0.96
jaxmd20 / HLOOpt / cuda / PostRev 0.005300387 s 0.005524066 s 0.96
jaxmd20 / HLOOpt / cuda / BothRev 0.002789458 s 0.002818 s 0.99
jaxmd20 / PartOpt / cuda / PreRev 0.002830803 s 0.002889521 s 0.98
jaxmd20 / PartOpt / cuda / PostRev 0.005477219 s 0.0056870079999999 s 0.96
jaxmd20 / PartOpt / cuda / BothRev 0.002766322 s 0.002860852 s 0.97
jaxmd20 / IPartOpt / cuda / PreRev 0.002807827 s 0.002924271 s 0.96
jaxmd20 / IPartOpt / cuda / PostRev 0.005434851 s 0.005666755 s 0.96
jaxmd20 / IPartOpt / cuda / BothRev 0.002798411 s 0.002837265 s 0.99
jaxmd20 / DefOpt / cuda / PreRev 0.0028287209999999 s 0.002912848 s 0.97
jaxmd20 / DefOpt / cuda / PostRev 0.002760339 s 0.002859119 s 0.97
jaxmd20 / DefOpt / cuda / BothRev 0.002771504 s 0.002833778 s 0.98
jaxmd20 / IDefOpt / cuda / PreRev 0.002819632 s 0.0028988649999999 s 0.97
jaxmd20 / IDefOpt / cuda / PostRev 0.002310646 s 0.002351861 s 0.98
jaxmd20 / IDefOpt / cuda / BothRev 0.002764978 s 0.002841298 s 0.97
jaxmd20 / JaXPipe / tpu / Primal 0.00927782875 s 0.009279790625 s 1.00
jaxmd20 / Jax / tpu / Primal 0.00927864625 s 0.009277066875 s 1.00
jaxmd20 / HLOOpt / tpu / Primal 0.0091559412499999 s 0.009156288125 s 1.00
jaxmd20 / PartOpt / tpu / Primal 0.009196729375 s 0.0091969675 s 1.00
jaxmd20 / IPartOpt / tpu / Primal 0.009198286875 s 0.009199265 s 1.00
jaxmd20 / DefOpt / tpu / Primal 0.008801955 s 0.00879701125 s 1.00
jaxmd20 / IDefOpt / tpu / Primal 0.008695781875 s 0.008693685625 s 1.00
jaxmd20 / JaXPipe / tpu / Forward 0.0174276625 s 0.01741432 s 1.00
jaxmd20 / Jax / tpu / Forward 0.0187232575 s 0.018728843125 s 1.00
jaxmd20 / HLOOpt / tpu / Forward 0.017388349375 s 0.017404439375 s 1.00
jaxmd20 / PartOpt / tpu / Forward 0.01741371625 s 0.01740671 s 1.00
jaxmd20 / IPartOpt / tpu / Forward 0.017410838125 s 0.0174164325 s 1.00
jaxmd20 / DefOpt / tpu / Forward 0.017415080625 s 0.017410795625 s 1.00
jaxmd20 / IDefOpt / tpu / Forward 0.01741167125 s 0.017411913125 s 1.00
jaxmd20 / JaXPipe / tpu / PreRev 0.025443836875 s 0.025446514375 s 1.00
jaxmd20 / JaXPipe / tpu / PostRev 0.021857848125 s 0.021859200625 s 1.00
jaxmd20 / JaXPipe / tpu / BothRev 0.0254474575 s 0.02544817125 s 1.00
jaxmd20 / Jax / tpu / BothRev 0.021856184375 s 0.0218556706249999 s 1.00
jaxmd20 / HLOOpt / tpu / PreRev 0.0255635125 s 0.025562918125 s 1.00
jaxmd20 / HLOOpt / tpu / PostRev 0.02072303875 s 0.0207051775 s 1.00
jaxmd20 / HLOOpt / tpu / BothRev 0.025665996875 s 0.02566237625 s 1.00
jaxmd20 / PartOpt / tpu / PreRev 0.025468950625 s 0.025475254375 s 1.00
jaxmd20 / PartOpt / tpu / PostRev 0.021510520625 s 0.02150513125 s 1.00
jaxmd20 / PartOpt / tpu / BothRev 0.0255508275 s 0.025568385 s 1.00
jaxmd20 / IPartOpt / tpu / PreRev 0.025451148125 s 0.025453468125 s 1.00
jaxmd20 / IPartOpt / tpu / PostRev 0.021516468125 s 0.02151529 s 1.00
jaxmd20 / IPartOpt / tpu / BothRev 0.025542141875 s 0.0255439025 s 1.00
jaxmd20 / DefOpt / tpu / PreRev 0.0254721175 s 0.02547305875 s 1.00
jaxmd20 / DefOpt / tpu / PostRev 0.018812451875 s 0.018810044375 s 1.00
jaxmd20 / DefOpt / tpu / BothRev 0.025548066875 s 0.025561076875 s 1.00
jaxmd20 / IDefOpt / tpu / PreRev 0.025454721875 s 0.025455626875 s 1.00
jaxmd20 / IDefOpt / tpu / PostRev 0.0183312075 s 0.0183291187499999 s 1.00
jaxmd20 / IDefOpt / tpu / BothRev 0.025542156875 s 0.025542485625 s 1.00
jaxmd40 / JaXPipe / cpu / Primal 0.075549823 s 0.089487879 s 0.84
jaxmd40 / Jax / cpu / Primal 0.078457626 s 0.0878360639999999 s 0.89
jaxmd40 / HLOOpt / cpu / Primal 0.106482933 s 0.112789086 s 0.94
jaxmd40 / PartOpt / cpu / Primal 0.070592279 s 0.080333378 s 0.88
jaxmd40 / IPartOpt / cpu / Primal 0.077138553 s 0.084480979 s 0.91
jaxmd40 / DefOpt / cpu / Primal 0.108419786 s 0.114900792 s 0.94
jaxmd40 / IDefOpt / cpu / Primal 0.0941775619999999 s 0.111910519 s 0.84
jaxmd40 / JaXPipe / cpu / Forward 0.182922384 s 0.201633292 s 0.91
jaxmd40 / Jax / cpu / Forward 0.098965478 s 0.108089939 s 0.92
jaxmd40 / HLOOpt / cpu / Forward 0.189008284 s 0.206739723 s 0.91
jaxmd40 / PartOpt / cpu / Forward 0.185237874 s 0.200147787 s 0.93
jaxmd40 / IPartOpt / cpu / Forward 0.185883699 s 0.201814104 s 0.92
jaxmd40 / DefOpt / cpu / Forward 0.184323312 s 0.20460292 s 0.90
jaxmd40 / IDefOpt / cpu / Forward 0.184844453 s 0.197395045 s 0.94
jaxmd40 / JaXPipe / cpu / PreRev 0.23758077 s 0.273711546 s 0.87
jaxmd40 / JaXPipe / cpu / PostRev 0.151025008 s 0.17284138 s 0.87
jaxmd40 / JaXPipe / cpu / BothRev 0.2523112 s 0.271164212 s 0.93
jaxmd40 / Jax / cpu / BothRev 0.149634422 s 0.168613278 s 0.89
jaxmd40 / HLOOpt / cpu / PreRev 0.247719974 s 0.264947864 s 0.93
jaxmd40 / HLOOpt / cpu / PostRev 0.207456636 s 0.227558183 s 0.91
jaxmd40 / HLOOpt / cpu / BothRev 0.281148006 s 0.296226275 s 0.95
jaxmd40 / PartOpt / cpu / PreRev 0.238805415 s 0.268889165 s 0.89
jaxmd40 / PartOpt / cpu / PostRev 0.150913794 s 0.1784504489999999 s 0.85
jaxmd40 / PartOpt / cpu / BothRev 0.259144872 s 0.314284117 s 0.82
jaxmd40 / IPartOpt / cpu / PreRev 0.245765283 s 0.258440415 s 0.95
jaxmd40 / IPartOpt / cpu / PostRev 0.159464735 s 0.157997828 s 1.01
jaxmd40 / IPartOpt / cpu / BothRev 0.273553767 s 0.301796951 s 0.91
jaxmd40 / DefOpt / cpu / PreRev 0.251284314 s 0.266499049 s 0.94
jaxmd40 / DefOpt / cpu / PostRev 0.216457377 s 0.223556878 s 0.97
jaxmd40 / DefOpt / cpu / BothRev 0.283092537 s 0.277931194 s 1.02
jaxmd40 / IDefOpt / cpu / PreRev 0.240946843 s 0.259998777 s 0.93
jaxmd40 / IDefOpt / cpu / PostRev 0.210905134 s 0.222898555 s 0.95
jaxmd40 / IDefOpt / cpu / BothRev 0.2666802 s 0.301388099 s 0.88
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / cuda / Primal 1.707896245 s 1.701157012 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / cuda / Primal 1.709686367 s 1.7031430600000002 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / cuda / Primal 1.722392512 s 1.714958551 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / cuda / Primal 1.702376352 s 1.69272958 s 1.01
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / cuda / Primal 1.699793791 s 1.6922887880000002 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / cuda / Primal 1.67041669 s 1.664518708 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / cuda / Primal 1.921296495 s 1.913615274 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / JaXPipe / tpu / Primal 3.03816736375 s 3.038750500625 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / Jax / tpu / Primal 3.038653240625 s 3.03936722125 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / HLOOpt / tpu / Primal 3.120886353125 s 3.121648468125 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / PartOpt / tpu / Primal 3.0594917125000003 s 3.060112531875 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IPartOpt / tpu / Primal 3.05965576125 s 3.06032794375 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / DefOpt / tpu / Primal 2.102160463125 s 2.10245704 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_24_outer_steps_4 / IDefOpt / tpu / Primal 2.94735541875 s 2.9484067300000003 s 1.00
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / JaXPipe / cpu / Primal 6.694297569000001 s 7.522609622 s 0.89
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / Jax / cpu / Primal 6.687179958 s 7.451253786 s 0.90
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / HLOOpt / cpu / Primal 6.649507911 s 7.300774703 s 0.91
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / PartOpt / cpu / Primal 6.806925611 s 7.5237576 s 0.90
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IPartOpt / cpu / Primal 6.760905395 s 7.461139831 s 0.91
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / DefOpt / cpu / Primal 2.722150853 s 3.231166343 s 0.84
neuralgcm_v1/deterministic_2_8_deg_inner_steps_2_outer_steps_2 / IDefOpt / cpu / Primal 7.488230872 s 7.826565502 s 0.96

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

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

Reviewers

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

@wsmoses wsmoses wsmoses approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@avik-pal @wsmoses