tokamax/__init__.py,sha256=qHXwAJC1YJVMjKwEava-P6rIzFTjJNJ71rhEzLk5MSM,2131
tokamax/autotuning.py,sha256=ozAUPHbs1_-t28etjF6pm1gKrBoXS65kziFLkQjXE6k,903
tokamax/benchmarking.py,sha256=sJxtRSAuHgW4JbGMyHw07vvUwRR_z8NRVYGVHCZrehw,1097
tokamax/config.py,sha256=FwB97rXW0P1VkMwhS-Mpti1LeBbdpC4l3AFcLAZXgKk,868
tokamax/tokamax_test.py,sha256=4Ui3s5e65BswYNpM-V_ZlBmrqtXvqNxmcVcxUM_RZL8,4296
tokamax/_src/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ad.py,sha256=gN-DnLTOWracL0cfa2dzq0Zc1Wqg7Z3-vY9ZzR926FY,4378
tokamax/_src/ad_test.py,sha256=_A0KujTTTdYH6XcdzhZ6KpoWYlKHWx911TbSavmS7qU,3782
tokamax/_src/batching.py,sha256=0hlm6Vnu-4SLDgE8qWEqVxWlug_gmG5ZMp6R-fZYdeM,7793
tokamax/_src/batching_test.py,sha256=_5y9wN3Vafu8X6oxpto_2Xj0yJzAfBMP0JU5t4WFDvo,4693
tokamax/_src/benchmarking.py,sha256=zKDhNrt_fGCm83oUgSvRsks033D1A5SA0eRNcbuogGM,17999
tokamax/_src/benchmarking_test.py,sha256=JeH1VMK-MYRhrAwAKSBTyDBJHvXc_uQ--dDQ84FN1i8,8145
tokamax/_src/config.py,sha256=i8uM8g9o0ef74sOrj8Sxq8_Yp6kdE-GWjkiMQDTYauY,2873
tokamax/_src/config_test.py,sha256=HuekTGgya8RPo2UYfrQfq8dM7FLonua7U6pUa825Vtk,1561
tokamax/_src/conftest.py,sha256=qjs1VrvMMZeefevbOkpj-TtXSqT2sZGJj-nQHH7Jrrg,1269
tokamax/_src/gpu_utils.py,sha256=8lNND9PmADJbz6ydY9Kg5d37yWYw2Dt4PYOgK7HDpto,3119
tokamax/_src/gpu_utils_test.py,sha256=5QMPfeky8zHTUwx6e9a2OBYC7T_olAiFAlj0GCxR5mY,1687
tokamax/_src/hlo_utils.py,sha256=Ox-9pNARzvhHViTPv4Eq2I28ggE8oItefb8BmjwqnQc,10138
tokamax/_src/hlo_utils_test.py,sha256=AuO8XmsCCytye_-H6EDxk5EOUL-pi3HZ_OKaLc-MeEg,12660
tokamax/_src/jaxtyping.py,sha256=_W4COs9cj3GdNE0B_J2QNhn86tt2JPTUaVzp8g_rtP4,1389
tokamax/_src/jaxtyping_test.py,sha256=vDDdV7JyVds-SpyReKiqZr1-WbAT-3qRnXloYLVRYGE,1558
tokamax/_src/mosaic_tpu.py,sha256=WGurHAedvkA-L2d_FhNRrQzRKLtjgflKd8h_4t25xCM,9303
tokamax/_src/numerics.py,sha256=QKk6nEq9L_XNSTn2AekAvwBceUtkNNo3U1a0VSlXBq0,8010
tokamax/_src/numerics_test.py,sha256=Inz-7R7aUwSiY6AGesixuWNABSUWO3OzVdryjFMjlpY,6204
tokamax/_src/precision.py,sha256=XRT401YjovNtpyeWwhiPVvSGHTLf9TvPHWuQpnoiWwY,7339
tokamax/_src/precision_test.py,sha256=fFifAxvC3-hv7tY4qLXUj6313laQkeIB9AN0c4uFxC8,3630
tokamax/_src/pydantic.py,sha256=VmLZX8h8B-e50hqMmPWsU3i_rT7xCsYfyZQdjY0Ehzs,10560
tokamax/_src/pydantic_test.py,sha256=BmvUxBQ7BDg03yDRRLqYfTZP7CYUNk7PKpVJpwZnFec,7726
tokamax/_src/quantization.py,sha256=gyqo8qBXUMcX5f15--v_zpflDjXqXqLnvA5V6WZ36KA,2913
tokamax/_src/shape.py,sha256=jw0gpUh0pg0T6UOw9CTCYXhIElmQwWx12GTr_bV7FAs,3036
tokamax/_src/shape_test.py,sha256=8C6awpma4oJzPlcIpqwVvf5SSfKmATfghDA6hZkr5Os,1560
tokamax/_src/test_utils.py,sha256=xh9tc0LnKp6SrT1VXY34HU-FmZzwh0CmbTh5_eRkt10,3205
tokamax/_src/test_utils_test.py,sha256=O1457xfEX1w4v5X8tW8s7SLCWNWZ9l_lxz5tausF_48,2097
tokamax/_src/utils.py,sha256=Hv8u5AztBwt8GK9YSn0PaDIlHoBZdXCKUxKg8zyZqqE,1941
tokamax/_src/autotuning/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/autotuning/api.py,sha256=gnsC538LUSGKHBRNHh-IZ5TdjD5EUGwlXQAMGyNUWis,11374
tokamax/_src/autotuning/api_test.py,sha256=w_TIOPizv1DOPD-onb6M4FYoKd3HVNdEplAQ2pc-sxs,13541
tokamax/_src/autotuning/arg_spec.py,sha256=-dL63Jot8KLDQYXW4EylBAsmCe0MyFlpsWEXnEI5paQ,1767
tokamax/_src/autotuning/autotuner.py,sha256=0SSJjfZigfiMUlwaFUK9rP8muObAD9qOoclvqx4i3rs,6910
tokamax/_src/autotuning/cache.py,sha256=qnNBkwxjFgBfWKVR7kPaNBV6EPUzrgUeT2GjWqJ499w,3425
tokamax/_src/autotuning/cache_test.py,sha256=OsXrSrsjHxelzl-wVmYo08lIbkPZYiI87UKiWuxHZo4,2171
tokamax/_src/ops/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/op.py,sha256=WvJSbT5NNlKjUY5GLgPQtN3IoeJrqrHKrXrSRRCU30g,23822
tokamax/_src/ops/op_test.py,sha256=V1F4GqSOlqADF4i7VOE9E_wT01BHjUz0h5qXbzTIses,8000
tokamax/_src/ops/attention/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/attention/api.py,sha256=6N2uuAeDgsyZAatyRfp8LMleBhPb8fdTky5YmESgwBI,7011
tokamax/_src/ops/attention/api_sharding_test.py,sha256=MljBPJQGQADiwHXAPxxfHMBWf_7mP2e6k3fK0R5FJpM,2984
tokamax/_src/ops/attention/api_test.py,sha256=sFuTNMFnihHCetroZQk6joVXynfmpwjZgVdKnGNQUvw,13907
tokamax/_src/ops/attention/arg_specs.py,sha256=F8_JuMoJ9LAgyhXAujGdR_bewHKWGLw-yUL1UW52njA,2902
tokamax/_src/ops/attention/base.py,sha256=9jybhnZM8wrog5E2YdYigeW3OR7hAu3TfLTbty2Gkpw,32838
tokamax/_src/ops/attention/base_test.py,sha256=L9gAlnSReP0oyNL-lpQ7i9yrXcyfMoPV2NilwdCYjCo,5286
tokamax/_src/ops/attention/bench.py,sha256=GpLexo7hy0O-s43tYemmd4VyRASlDqpwvoEgmbKCAfU,2622
tokamax/_src/ops/attention/jax_nn.py,sha256=IK-ZxJO3jGhJeoypJL50Apw076fF5j1gOmS0VdZMTIM,5746
tokamax/_src/ops/attention/jax_nn_test.py,sha256=LXUF98jnagpQPRDg8cIRllYTLN78W3LGG4EdeT_Xb00,3889
tokamax/_src/ops/attention/pallas_mosaic_gpu.py,sha256=CDezVQvxmLlfy7Yw9Oct3Xe1IwdQaZxaUTGyGDKBndA,10856
tokamax/_src/ops/attention/pallas_mosaic_gpu_common.py,sha256=sjrrrC4UqOCmHePmZwcufybNvTfY3tLWwOvHODjjWdk,2667
tokamax/_src/ops/attention/pallas_mosaic_gpu_kernel_sm90.py,sha256=wh3bv5GFDc8QpIQ9tlHUbAKUyVHxaMJeypjL5whCrLU,17078
tokamax/_src/ops/attention/pallas_mosaic_gpu_test.py,sha256=gXlsg0YxSZO2FqyITEs_3fGrfB1HZIQxmZARCFMp_RY,6917
tokamax/_src/ops/attention/pallas_mosaic_gpu_vjp.py,sha256=xmnvRV9dPK4_O1xE0aN9FE9zEsXbxFtWHmuODrInoCE,7696
tokamax/_src/ops/attention/pallas_mosaic_gpu_vjp_common.py,sha256=eqxSg-ScoSs-AnDX-eN9-ttBNGXn5HO6uruaVeJtkDM,1114
tokamax/_src/ops/attention/pallas_mosaic_gpu_vjp_kernel_sm90.py,sha256=DSCd-vw2ejXeHqzKaEkfGccadgix-56nfD_8bV8IDmg,22406
tokamax/_src/ops/attention/pallas_mosaic_tpu.py,sha256=Vc2N4xpIgWqO_pD48vdSiynrkXamp2AWeFZA0IFuB4A,8390
tokamax/_src/ops/attention/pallas_mosaic_tpu_test.py,sha256=HLneEdBq8nD5xpcY9qHzk7CbGAuIDig-pwY8qSedwcU,3709
tokamax/_src/ops/attention/pallas_triton.py,sha256=c0e-l-E_ey975hsbmkdC0ryYzvrNIexLCkJ9oz_vGaY,23720
tokamax/_src/ops/attention/pallas_triton_test.py,sha256=MWV1W3tprgO3-EGOlSgrhB0HQcp6IOiWI0Y1LhZeC5k,5752
tokamax/_src/ops/attention/pallas_triton_vjp.py,sha256=X3Ir4NKhpJl6RQt1C6dx9H16q_yo2RLUjtYgjmWfhSA,18274
tokamax/_src/ops/attention/test_base.py,sha256=c4WvIk4nJ_mK-PJqvOWQxYKWzT4C_xoK84A78ehrlxo,38981
tokamax/_src/ops/attention/xla_chunked.py,sha256=y_nYk0Zc0Enc_mY9I1nJjYHzsApL5NGp9O6uSxIsMx4,12483
tokamax/_src/ops/attention/xla_chunked_test.py,sha256=lkw05nDSqgtsO0Ahx2W5i2i8A4HSmW51Y5lnBa1IfbM,8794
tokamax/_src/ops/experimental/tpu/splash_attention/base.py,sha256=kxeaK81f17OhzKGFuzNZ_7_3zdLFF3fRcrSEEDt_s-U,9058
tokamax/_src/ops/experimental/tpu/splash_attention/ring_attention_kernel.py,sha256=1HWH0lg3r1Gz497xPnCOHXKAigyCgxVR9pHvQndC0Y4,21284
tokamax/_src/ops/experimental/tpu/splash_attention/ring_attention_kernel_test.py,sha256=vcVCzjbwh-TbGdVmzDaQy987TEW0of0DYjXC9c0d8OQ,9325
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel.py,sha256=XgKP0LHIqPLRjSzeJIu-F-C5bJ9DVwHj_M_L2r4FHAI,67015
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel_sharded_test.py,sha256=fGIDDwRnpWKxLBAwIAIpC6-t5UgzaL-voCztFltOvc0,7235
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel_test.py,sha256=kQFPkD-jrmCR-iiFj2Z4Fjc7R9v2cTQsVdUDr59P5l8,20089
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask.py,sha256=oAYV0-2SHK6d4heY191G7s3QzJzLw96gF9iIUXfAwtk,14948
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask_info.py,sha256=s7GaZL3Lfi28er5TUMgGs_zoQN1wMNDTJ0_iV_yJkpU,22739
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask_test.py,sha256=30EDDtqAkqJdNBnBpMxyQ2qz90Bo9a0MzhSWa9rXuIM,50905
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_test_utils.py,sha256=2eyOdAbTzCKt20qnCvSTCk1da0mxCBn8Rh9jCOQh8-M,2279
tokamax/_src/ops/flex_attention/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/flex_attention/api.py,sha256=Td5kwvuksLSyCRlPpFCAjkpZkryRmywjpoRc3_1p9bE,1131
tokamax/_src/ops/flex_attention/arg_specs.py,sha256=DECEEmyM-WzMoD7z8nd48JBDKUK9zyhY2InIT76AVaU,820
tokamax/_src/ops/flex_attention/base.py,sha256=YBCBaDpAmIYuNSAbv0WpJevJfykmAvGH65kJril-a5s,12199
tokamax/_src/ops/flex_attention/base_test.py,sha256=l4y8geGYw4AIJ_HnNOfXIyHXw2LKVc1xYEBlwNhIws0,1475
tokamax/_src/ops/flex_attention/pallas_triton.py,sha256=PXehpvH4twjfoea8dwVXHPx1xDrbpGio_c03Mnz3-lg,18437
tokamax/_src/ops/flex_attention/pallas_triton_test.py,sha256=Nau27kj4rnaZunL7gOA5x06FlbNr6kBrICz-hNnbtZc,1896
tokamax/_src/ops/flex_attention/test_base.py,sha256=cw8HWDq1pFWlZ0qYoe1VgTPHyHileXKtBKEE1ZONGeM,2959
tokamax/_src/ops/flex_attention/wrapper.py,sha256=n0PHAmV8IPXG6LmH_Uh5PWN6UOIuazzWRQIYeS_DrD0,3481
tokamax/_src/ops/flex_attention/wrapper_test_base.py,sha256=yMK2Hm8M2QyCSkf2x2iFRlePbb5WQxR3bZhV8M9fWG0,1335
tokamax/_src/ops/gated_linear_unit/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/gated_linear_unit/api.py,sha256=fMsgesSy5-YrPfkXj-IOAKsMYXbPM9ONesD2uHvyOO8,4375
tokamax/_src/ops/gated_linear_unit/api_test.py,sha256=YCfwcxKB-nsm5hHV50thEKyctsYkbJszB9ItrkmSlLM,3696
tokamax/_src/ops/gated_linear_unit/arg_specs.py,sha256=DECEEmyM-WzMoD7z8nd48JBDKUK9zyhY2InIT76AVaU,820
tokamax/_src/ops/gated_linear_unit/base.py,sha256=xo_fxsku-Cwvj13xm7k6n7yWgZTqy5pBUnEooCZuzxc,6542
tokamax/_src/ops/gated_linear_unit/base_test.py,sha256=zNoOZbucSKmx2up3WFSuuiNG3Nw8wBzs3ekHfm21vUw,2701
tokamax/_src/ops/gated_linear_unit/pallas_triton.py,sha256=XSnDwLLgD8L2HaLt0Mluh9t8d7HPgEb4qwXXrESqOwE,8384
tokamax/_src/ops/gated_linear_unit/pallas_triton_test.py,sha256=0CYyz24N5KnqqpGl4zdvWmtEjU9GAA4fP2m8SqSylwk,1692
tokamax/_src/ops/gated_linear_unit/test_base.py,sha256=tA4fQQoQydeXgcmE6NnA8Ku901A6DE7e7H6DbERoLTE,3702
tokamax/_src/ops/linear_softmax_cross_entropy_loss/__init__.py,sha256=KjQvwjX9P90Ur6p00QPLGux-pBqPQrSYlirJoIWvhoA,732
tokamax/_src/ops/linear_softmax_cross_entropy_loss/api.py,sha256=8SUBww8T1y5Uc4Pn-k2WLXjB0ND-okiuc8UsQy0Uwmc,4343
tokamax/_src/ops/linear_softmax_cross_entropy_loss/api_test.py,sha256=o0_epX0qM6YkwNAkkGKYnAL1xdlYGtAAbmaRPioBOzw,4443
tokamax/_src/ops/linear_softmax_cross_entropy_loss/arg_specs.py,sha256=uDdT56_VKPdA-oj-sfk9ihCYQXgat1Vuoh25Vh51mWk,3841
tokamax/_src/ops/linear_softmax_cross_entropy_loss/base.py,sha256=IXs5poYvF2UgW9VUVrr194Qc2MkjXb0324_Rko-bpGM,2810
tokamax/_src/ops/linear_softmax_cross_entropy_loss/base_test.py,sha256=PlmNj8iET7fNEmxpAgTUtadowNwiqCG-8vkZT0N6DtQ,2494
tokamax/_src/ops/linear_softmax_cross_entropy_loss/pallas_mosaic_tpu.py,sha256=wwiej4NRtzvj0MxkIbyDWuAbGpbibr_MjtlmzJoSVto,23708
tokamax/_src/ops/linear_softmax_cross_entropy_loss/pallas_mosaic_tpu_test.py,sha256=YCBkT4Ln51lEoTkXAkAl7q_8FKYvDAikTgZYRXQx9zY,10453
tokamax/_src/ops/linear_softmax_cross_entropy_loss/reference.py,sha256=I_z_vvq7lFr7llMzaIXGUffmVKmdp5YZUmHBwuewyXE,2505
tokamax/_src/ops/linear_softmax_cross_entropy_loss/reference_test.py,sha256=mMTXhvM0A7tx8Def7jPzoVqOz4NF5mQalA-GC_wkav8,3703
tokamax/_src/ops/linear_softmax_cross_entropy_loss/test_utils.py,sha256=8kffMQWUuqoUvRTRO6x1mZ0d07LrJtriJmp_ehnB-mA,1679
tokamax/_src/ops/normalization/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/normalization/api.py,sha256=2tMls0EN_Ko5PunqVOMs77-XZDlIO-oZzxQ5akGLE_c,4609
tokamax/_src/ops/normalization/api_test.py,sha256=zTEs9gCN1tkotKYNBLGYp_U3KWcx_xVpEqgo-szPG2w,5368
tokamax/_src/ops/normalization/arg_specs.py,sha256=trQYL9acYBYVBO16mdozgBeHbWSUSlDuKljK6UZ4nAs,2739
tokamax/_src/ops/normalization/base.py,sha256=9oMy73lUaVwMkkwXF2nQGr1HjbXXaBFG6qSgAd8XrrI,6320
tokamax/_src/ops/normalization/base_test.py,sha256=XKZTDAMwsol5oEzzdB9Wu-pu2_12eSSadGLyMeKKweg,1027
tokamax/_src/ops/normalization/bench.py,sha256=cME1KGCi7yxlepZd2YH6zGZOz5wVvfbi9bePk5Kr3Jw,2193
tokamax/_src/ops/normalization/pallas_triton.py,sha256=3qL2HT2JMSXYOfbXXRco50R4tQ7-Anh9iqmO0vU5DGo,7993
tokamax/_src/ops/normalization/pallas_triton_config.py,sha256=etHhS3vixlLTOTttFCFxe4ro0KKsyplDBZFKqPi09Bs,4112
tokamax/_src/ops/normalization/pallas_triton_test.py,sha256=-s9elX36iQbgaFIlVK5Ii1qWSYzHWeX6SesWDb6sMBY,5154
tokamax/_src/ops/normalization/pallas_triton_vjp.py,sha256=JbsM0HyFOub6BlgLKKe-ICQ3WPMHcm0TIu5f7I9-QqA,7283
tokamax/_src/ops/normalization/pallas_triton_vjp_config.py,sha256=4EUNT-bFybyWiWz8GLa1iZwWw55eOaREmh345Plnlps,3118
tokamax/_src/ops/normalization/test_base.py,sha256=shzt2qWy3HsdgQqdC1vSEcOoF5-xQs6nsjQO4-qlTb8,6683
tokamax/_src/ops/ragged_dot/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/ragged_dot/api.py,sha256=bwZeFbugvMJNP6LuZc5IRQ7NRSTHTa6kbgcDrXY_s98,8435
tokamax/_src/ops/ragged_dot/api_test.py,sha256=hNeQgOXMSNRSCtbGtbmeRd_KzjdY1F2-uiN7aUzuhPQ,7405
tokamax/_src/ops/ragged_dot/arg_specs.py,sha256=rkOiyK-tmCIbzRYmoJBBfEBhggDDi8i29HcK61AUpZE,5013
tokamax/_src/ops/ragged_dot/base.py,sha256=TGf5ADGM06XEkhOsbadweJ_uXenKsQ35F8GLgBFrwr4,10669
tokamax/_src/ops/ragged_dot/base_test.py,sha256=Q8VIIfo7boCCMZ7-4BNNnn09VEV77w5DPEpy2Kmq73w,1525
tokamax/_src/ops/ragged_dot/bench.py,sha256=TaACBBxj0QpE9NcK3MW6JQL-ccqIiS4ZBmHyxRjIESs,4097
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu.py,sha256=abxAaoYHD_ulqBOgsuXuQJtuhoLdscBny6mf9bFDI94,12330
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_common.py,sha256=1eXs7-gjW6NgPHdkcoKxjYfsije48rr97DlGTit48Ro,14207
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_kernel_sm100.py,sha256=avStTTO1A2GphtYiO0OjQ9vtFQec4CDk-Pdm5gzwmvA,10879
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_kernel_sm100_quant.py,sha256=7HdADearqV0EOA4mRe9PqUKL_2CaicSKhZP2KqHwO3A,17006
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_kernel_sm100_quant_post_scale.py,sha256=2W9Uo-PbHFwiYOmnktlKC14ulQcvBthBUsy4hpS5x78,15841
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_kernel_sm90.py,sha256=drkhjwB7QLQ3wPJN3iFZYV1wm-5l-WuecZ9ybwi_-5c,9447
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_kernel_sm90_quant.py,sha256=u3qGFwdaLQEwwWcwq7EcDIkhvUpMYqa6ZW5GExT8y_o,5436
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_kernel_sm90_quant_ws.py,sha256=y0o0VTxv4B5mhe9JMcndkyb6gZusxnbl87Ns3MPFnSM,6241
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_kernel_sm90_quant_ws_async_store.py,sha256=kK9WpYqNyK4_tyTzV7M6Tcpx1kbWReMV5WwPCgr3ugo,14421
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_test.py,sha256=F_VjsW6x1H4EXeyI7u5zGsDS_yrtHAR4SXpsRaxKscw,2960
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu.py,sha256=HfZGYJpCxpddJYH4CUVU-G-MS9gqNv8QC_Wjnj8vEy4,16706
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_kernel.py,sha256=Cb9UJFa1N9TLWgZPfrVAEQcXlR14GEJjH4-KUumZ0SU,32696
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_test.py,sha256=3fMXVNg9lf2hCU03drSW_WdAy16eRXYWXBuzA_DlwxI,5464
tokamax/_src/ops/ragged_dot/pallas_triton.py,sha256=QDcookfR7XNgxhDSy91rLqhhqH2aJWKgvMUTLMQ26K4,15513
tokamax/_src/ops/ragged_dot/pallas_triton_test.py,sha256=Z4jwcJ4ou9i8cD2UKIet7Lq_8HnqU1IBkVhWEWl0OwU,2427
tokamax/_src/ops/ragged_dot/test_base.py,sha256=RSaWBjCslQN5W9BbxYqahHE-VOCmj-5yypMNg6HpDFA,10026
tokamax/_src/ops/triangle_multiplication/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/triangle_multiplication/api.py,sha256=55rXRwL-y1TrfsznoDx84TRt3UAa7Kw_idxCiMZCNpM,4360
tokamax/_src/ops/triangle_multiplication/api_test.py,sha256=h5s0pMZyElSv49Wx_TPEi6VsQXCtN2j1nMxMIAFhxSw,2894
tokamax/_src/ops/triangle_multiplication/arg_specs.py,sha256=BN9SSgfaCd-2Bz_CQMGbUQADuRx2iMn9heeWPWe-WM4,1931
tokamax/_src/ops/triangle_multiplication/base.py,sha256=fq_nwnV1aueOvuslRAydaLu-6JWOcfPvm_DcIILyEAQ,6446
tokamax/_src/ops/triangle_multiplication/base_test.py,sha256=LLjXY0wTXuUMsradIVGN1VLCMI3zhj145saHr4XeZbU,3452
tokamax/_src/pallas/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/pallas/block.py,sha256=P2LjA-rm5uRCpYY4dUTdKVCkDAvfYh5QhmTYwh5IOTo,11529
tokamax/_src/pallas/block_test.py,sha256=s1lNHBl_guAh2BwMjjWiqeKR0Nytl8jn9SF0dkFdN5c,5519
tokamax/_src/pallas/grid.py,sha256=d_T5YafLTkixhDr8otTNh5frmIcpuHosYrWTbCnqAKQ,3641
tokamax/benchmarks/attention.py,sha256=QkN-Lgnsjkcyy8Yzd_Em76QZ331weMVXjrujhBGAVow,3356
tokamax/benchmarks/triangle_multiplication.py,sha256=-pZbLF1P7WQXWyXMNLUMWXDN6fw_Qn1AAvNLtIZ5TLo,4097
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/dot_product_attention.json,sha256=mBTg46sLvlh-LCUBsZ2VyGFvmr3PU-mxDeL91FascIA,3665
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/dot_product_attention_vjp.json,sha256=zCD5h5Z7P47Bg9OvmQLH4iSk-o8FHYXMU1zKS881ebY,4237
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/flex_attention.json,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/gated_linear_unit.json,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/jax_nn_dot_product_attention.json,sha256=oyJNe7PV79f-JGn31hb8T0HBSQXO_8PP2oBq8nCxYOc,4279
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/linear_softmax_cross_entropy_loss.json,sha256=-QoPczaaIQ52tQDQpGMJlCv7jNdC7-9R-_wETIewxUQ,39807
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/linear_softmax_cross_entropy_loss_vjp.json,sha256=LDb3jX-sOtgGlk6KvOgfW7cXQdFONZKqvJGonjQOxWU,46499
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/normalization.json,sha256=2KeLHQKNKTYxp4oFx4aFJlvA1J8v4HZ6SfAvL50XHwQ,2466
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/normalization_vjp.json,sha256=nh3drbifL0sc9GU3IvVyacoHsRbOQ-LejzCRjJ0Si_w,2448
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_mosaic_gpu_flash_attention.json,sha256=OKAYjRaQBF6WGgi__p-8ZaK2l1BwS3S9s3-vtpJktC0,4476
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_mosaic_gpu_flash_attention_vjp.json,sha256=2xorg6pYgrKtx1qSZrKMWpFBYQz4hxKL4Glv5rB5kZ4,5451
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_mosaic_gpu_ragged_dot.json,sha256=D6LpRpKK1dFjpd9HIN7cJtGlXykgfOh2tHaCVaGSY5k,9043
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_triton_flash_attention.json,sha256=aadm488fG0RoyG1WuFJWcF4lGrFPBAQN-ps5VolJytM,4888
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_triton_flash_attention_vjp.json,sha256=aOoWJl0kYpHQiX1vFZ-v9Q2iCXfIEpqbqRCA7bIvbmU,5401
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_triton_gated_linear_unit.json,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_triton_normalization.json,sha256=kvKhWPajdF0uD9oXplGDV_Tl1zAsahCclU5dOcQWWk0,979
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_triton_normalization_vjp.json,sha256=iH6HidLYZ-sU3HO3b0TNfXoMAZ0gJfrVWHmHarLdQIE,3351
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/pallas_triton_ragged_dot.json,sha256=2EqPg4ZoYnHiMcSuEvzkzLrDMuU-9gSBN2uo2aWsXLA,8061
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/ragged_dot.json,sha256=Ap_4QfC3rd3GgRYtPmN4vEeDFvn6P9oIzESUJjSjLJ0,6168
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/triangle_multiplication.json,sha256=vt8V7ATaZnYAmEPy_7xcDO7jRaT22eMBb7a3RotxdA0,6464
tokamax/data/autotuning/nvidia_h100_80gb_hbm3/xla_chunked_dot_product_attention.json,sha256=HZ0T15lOTKNiUehM9CPCRM7ubMn_1w7al34sLs8B8LU,4278
tokamax/data/autotuning/tpu7x/dot_product_attention.json,sha256=2lySXjGoUuyS_0fWU0bRlNxctfEEMIckg-YEz9usxdU,4148
tokamax/data/autotuning/tpu7x/dot_product_attention_vjp.json,sha256=lSWcB_X14gGimvXEtVQoUfmKePtGzQQHJ_jA-4iheC8,4501
tokamax/data/autotuning/tpu7x/flex_attention.json,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
tokamax/data/autotuning/tpu7x/gated_linear_unit.json,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
tokamax/data/autotuning/tpu7x/jax_nn_dot_product_attention.json,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
tokamax/data/autotuning/tpu7x/linear_softmax_cross_entropy_loss.json,sha256=3P7OkGs-N_1Vdbxnk04WAXXA8rqPOGVwti34IvLBTaU,42454
tokamax/data/autotuning/tpu7x/linear_softmax_cross_entropy_loss_vjp.json,sha256=JxlY8X-ez0R9JXDleH6ERhyVOA29PzpfOjXJ4KSzli4,49699
tokamax/data/autotuning/tpu7x/normalization.json,sha256=YgkVwnTJXQ7ZF2FHNpPEeKHzKhRVPMmqMjjBYF0yqV8,2378
tokamax/data/autotuning/tpu7x/normalization_vjp.json,sha256=thpKZxR3iBHDIcIGmTah7D3lu20rJ84ITt-NQVB8TxI,2387
tokamax/data/autotuning/tpu7x/pallas_mosaic_tpu_linear_softmax_cross_entropy_loss.json,sha256=ROGPld94cFEgheHRfdqJOvm1btLwlSdhJ54uitfQxxI,25283
tokamax/data/autotuning/tpu7x/pallas_mosaic_tpu_linear_softmax_cross_entropy_loss_vjp.json,sha256=3nmlDWjEeiYn5J8shy4RlSkWSCbm7M0Sm7Yq53bawtk,32325
tokamax/data/autotuning/tpu7x/pallas_mosaic_tpu_ragged_dot.json,sha256=hFeRDL-eW3gSy-fYYdopz1eIy7t7TiHDyhEUnl2Gouo,7527
tokamax/data/autotuning/tpu7x/ragged_dot.json,sha256=3a7Ag2DuLtWuGKFh_bvr83RZljmYQ0ZZgO9XuYnArqU,7404
tokamax/data/autotuning/tpu7x/triangle_multiplication.json,sha256=pLv1_WFDk8AmFe-PR9JSZms6FX-Ka0YFx8wKjQnIY1c,6264
tokamax/data/autotuning/tpu7x/xla_chunked_dot_product_attention.json,sha256=joJS0kSpOkVn_AHUr-oj6ntQuUJA8yWr_fPqFISmnxo,4170
tokamax/experimental/utils/tuning/tpu/ragged_dot_benchmarking.py,sha256=KkpQ7jM68IPDhMKx6CISeATeOT4veMLVGvMo83ktbsw,12484
tokamax/experimental/utils/tuning/tpu/splash_attention_benchmarking.py,sha256=3njWkUe3I_lGqO25EruH4kQmY-HFCRQPkBE2ndLQmYE,10382
tokamax-0.0.10.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
tokamax-0.0.10.dist-info/METADATA,sha256=P5Z-gOkvMmVnbJAwFdt2yGhW2v0XlZvp8T6ztuoCUD4,21125
tokamax-0.0.10.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
tokamax-0.0.10.dist-info/top_level.txt,sha256=xHgO7qOuYjrLH3047bLQLnwrjI76i9-xlChxNQIZn-Q,8
tokamax-0.0.10.dist-info/RECORD,,
