tokamax/__init__.py,sha256=EEZ0twht-1FRD0vvRineeDFDxpStk7qQ3kh9wkM6I_A,1717
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=lPTCs6bfK3-D0rOHhFLsxUhWNk3uQ6JIcBmN33MUsUs,3903
tokamax/_src/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ad.py,sha256=gN-DnLTOWracL0cfa2dzq0Zc1Wqg7Z3-vY9ZzR926FY,4378
tokamax/_src/ad_test.py,sha256=5e-2DkE4i7qE-GBpoajwvUFuQAoBRrOWeids0AaYl2c,3701
tokamax/_src/batching.py,sha256=z-OS9VTui0ahhaLZGu_oZrRphgTekXnDw0tHfWTHbW4,7974
tokamax/_src/batching_test.py,sha256=voxWU-AjoDFXIAoSeI6HA9peNntiP_2jueETvKh959Y,4127
tokamax/_src/benchmarking.py,sha256=aKglc2hz0Un66GD8DRzhJOIh2QaGHrqoRu1m3t1FOPA,17473
tokamax/_src/benchmarking_test.py,sha256=2ZqhdUWh2AgXUyWbw0JdzAeltTW6UDQZTww92wQDt6Y,8016
tokamax/_src/config.py,sha256=1RDPBF4Uahsoi2rR4TyyT1sk0pqUNWrxee1NlcBMCfk,2786
tokamax/_src/config_test.py,sha256=HuekTGgya8RPo2UYfrQfq8dM7FLonua7U6pUa825Vtk,1561
tokamax/_src/conftest.py,sha256=PC6zlS73bx_BB5pFvbOTB-ZASHZojUFrUReMVtozP0o,1225
tokamax/_src/hlo_utils.py,sha256=0xxBUVsbydCj23h9qevpDNPtuyL7XXDRhi-hro0dpLk,12995
tokamax/_src/hlo_utils_test.py,sha256=ZFwDBpuyHsyRyP3b1SKGfBa5f2Ixf0ePnaOEyBYUrm4,16819
tokamax/_src/jaxtyping.py,sha256=_W4COs9cj3GdNE0B_J2QNhn86tt2JPTUaVzp8g_rtP4,1389
tokamax/_src/jaxtyping_test.py,sha256=vDDdV7JyVds-SpyReKiqZr1-WbAT-3qRnXloYLVRYGE,1558
tokamax/_src/mosaic_gpu.py,sha256=9J4gDBXXIvGQPcssGZhXF2Ci0hDOjYX66VcoNju0IU8,1037
tokamax/_src/mosaic_tpu.py,sha256=IFzDkbtAeExVD32vaWsJPer5paUOQACP9JtlYAKEZk0,7833
tokamax/_src/numerics.py,sha256=LwB4I-XcFNSd8ILUMaL5EHmtXvL54RI3ZXtBNdBW1Ls,6246
tokamax/_src/numerics_test.py,sha256=VWRgdU3S-Bhgj3y6DKVx2xuvhAJh2uJrip8z7Zl1nFQ,5470
tokamax/_src/precision.py,sha256=VOsGoH52xxuxN_7YN5xHPFhqC13ILp7uBs_jSD7DL4E,6293
tokamax/_src/precision_test.py,sha256=Cpba1dPi8mKK22gfb4fCyAd1NE62MAd11Du3GUx5DRc,3584
tokamax/_src/pydantic.py,sha256=DwX7yrOp17VCcWMhGPDQlMFwODoR7T16U8YTuUKnY60,10168
tokamax/_src/pydantic_test.py,sha256=Akn8HWbzCwM66RAIajip-iF6zbt8ySA7Nz1EPLHEUlc,7758
tokamax/_src/quantization.py,sha256=iSJmsT8bFSJzt4NTluT4V6mgaMju4CPJ4GNKRbukSN4,4538
tokamax/_src/quantization_test.py,sha256=cFEJrvZjat4TEjz2kERRxdA4J4C3fRVqtVbwbmn83TA,3133
tokamax/_src/shape.py,sha256=TCdgVWEcP-cN0vrKYdSdJcBR2H0nremenanWbyEZDe0,2103
tokamax/_src/shape_test.py,sha256=QCqA-EEqG1cUxM-xCx0jFgqotEaQjrV0nXIPVTfZTV4,1228
tokamax/_src/test_utils.py,sha256=QKNNgaGOBMNbaO-IUL1gbXydachSMhmIW2URPiRtoSE,2023
tokamax/_src/triton.py,sha256=CSCgX6rw5UcnlQ3uJdnr1luFWuyGljXDWPkubyNgnw4,1369
tokamax/_src/triton_test.py,sha256=0Y-4nft57BwZvII3JZJTu6BjdCu893zj7D9hrcVAmhM,1045
tokamax/_src/utils.py,sha256=3bYppf7Qy9QXtar_Bw8K5L31bx_MXHyC_aC-Is2ynj4,1536
tokamax/_src/autotuning/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/autotuning/api.py,sha256=TjxQ2nIwc7veegnGAsvN9eoOPP0KmbxqDjc1DXLFcEk,10545
tokamax/_src/autotuning/api_test.py,sha256=MS2_IMHKO_bh6JvrQhmy9hE0FGJUSHkLhHOLYtoDRsY,12857
tokamax/_src/autotuning/arg_specs_common.py,sha256=zwU2OmwwyolWenORsK-Uvkg_mDzp_IYzzkr94Yfw1M4,1896
tokamax/_src/autotuning/autotuner.py,sha256=YHAmOb35Xymd50kGoTuf1NcMAGYjxy9FD3YgEyn4cyo,6751
tokamax/_src/autotuning/cache.py,sha256=xac1Ct-jOE7dhUwaKa-XY-2j2sfnpKDFWAbSMGxhz2o,1564
tokamax/_src/ops/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/op.py,sha256=f-PLXqQPety8GIhTQKpLS03QJfAMY3I8yLdwharVgIU,21249
tokamax/_src/ops/op_test.py,sha256=ePxgS1nYfTMosH2hdogYQ_fTEJgj9x13dy6D8yLdumU,7644
tokamax/_src/ops/attention/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/attention/api.py,sha256=_Tkn7fojxqQ643p54MVuRTeaBYQtXGxmLnAjTkD340I,5852
tokamax/_src/ops/attention/api_test.py,sha256=clHl_x7DCu3orjV6P4mlr4reZAFTxq6VFzYnChoZdW4,12079
tokamax/_src/ops/attention/autotune_arg_specs.py,sha256=IvLBZb8qbl5uZ9Ah_a1KeeAfx-e_MNRkCsGbDaYmjH0,3471
tokamax/_src/ops/attention/base.py,sha256=u3fQSRKRa_hO9U2cCq1ql02HD6TRmX4oyqooQ6LZXx0,28809
tokamax/_src/ops/attention/base_test.py,sha256=L9gAlnSReP0oyNL-lpQ7i9yrXcyfMoPV2NilwdCYjCo,5286
tokamax/_src/ops/attention/bench.py,sha256=D5hjAFCud9VoT7Twpn0xgfxazLsCQubj9S0puGk_aEQ,2660
tokamax/_src/ops/attention/bench_arg_specs.py,sha256=gIfT4JdUXtaFfyMe7bOafHI6yuNzuZWFGOBV0lwwnCk,1699
tokamax/_src/ops/attention/jax_nn.py,sha256=qUVr_C9habnOvfXanYVx1U2R0JmcWmONhZ-_Nt-GTuo,5532
tokamax/_src/ops/attention/jax_nn_test.py,sha256=WNRmKjHjYrRIFZfOm74c849My7pAqPeKvGdaiqrI0-8,3809
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention.py,sha256=soZkqJkzioPir8AX1J93fwU3ZTSIsrvRUaRJeaf9-z8,23926
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention_test.py,sha256=aSMc-jCaTcE_qXZxiZT_JsNVmFvtIJRw-HtDr0ZA1Yg,4812
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention_vjp.py,sha256=RYnDxlKUdfHN0SMoFzIWK5gyw8Ufkfp7FbPIZjJn7ss,29949
tokamax/_src/ops/attention/pallas_triton_flash_attention.py,sha256=ns-Mr79eQE77Pjcu5q7xZbdlE6TsgEoSCDbnMLdo5N8,23886
tokamax/_src/ops/attention/pallas_triton_flash_attention_test.py,sha256=rY-ZsdAZieth0p0wPLlmMWBXzu9k88TGk8fnCqvpKr8,5732
tokamax/_src/ops/attention/pallas_triton_flash_attention_vjp.py,sha256=nd5NqMldmMKaGq00k69uYUQ3twN1qtXLd0842hgQByI,17482
tokamax/_src/ops/attention/test_base.py,sha256=z5xEI7OYS8aL1UELn-cl8y_BpIDZS8eJJ72lfzo3c5A,36319
tokamax/_src/ops/attention/xla_chunked.py,sha256=plWS_LSinFCtrAFTIHoM_MrByK64aYNtHzOJlaD9Jx0,11525
tokamax/_src/ops/attention/xla_chunked_test.py,sha256=_y92iukPWgv5Yrdu5yE5lLTWKm4ErSMlVT_vKiDyd3A,8984
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel.py,sha256=dsy-3iehiueUFDwFSDkMh0ixnunuDm51Z0K7dd7lOhk,71805
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel_sharded_test.py,sha256=S8OTVQ9FPQ6T9Lzp4bVMrBTvpC6-5HHsYarl-qyWxWE,7266
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel_test.py,sha256=zzqlV6R2t1BbR8Sj0jpPLivNf-ImbGrNf7AI6lV7MS4,22222
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask.py,sha256=uz8Mc2hQlp3_5LUQlkT4tkrnx3SQfeXopk5WaLTt1tc,14938
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask_info.py,sha256=hFWmTVtOzHt8vtf-0zwuYVjeBpm-4lPHbaT80qzRXvo,19501
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask_test.py,sha256=V8ze4bpcmR4r_jX9D16niVYmGkTphKp0I0C_imQr5zY,45801
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_test_utils.py,sha256=mt7VggyVVlhUfaW4mVYA0_O-_R2TvjRb17ETLU3fFp0,2278
tokamax/_src/ops/flex_attention/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/flex_attention/autotune_arg_specs.py,sha256=XEUVNtsXI_ykm5wW2QX833NxmY55o8uIupFeS8BfJXM,831
tokamax/_src/ops/flex_attention/base.py,sha256=02HBWZLyEbBjrinOoDINDtHyYELULmVFPE71P4ZDR6Y,12222
tokamax/_src/ops/flex_attention/base_test.py,sha256=l4y8geGYw4AIJ_HnNOfXIyHXw2LKVc1xYEBlwNhIws0,1475
tokamax/_src/ops/flex_attention/pallas_triton.py,sha256=qU6QxGCi2S09VyNLhPNOqjWPz5VjeMhkhLRhG5-eCC0,18143
tokamax/_src/ops/flex_attention/pallas_triton_test.py,sha256=49RoyGMHq5y_D0yCwTDXdTQY3--IfE2I3Tp9NFiQAto,1887
tokamax/_src/ops/flex_attention/test_base.py,sha256=GCzJ_3jI_26qGX5lYVEKXodeFdyc9mNkgeojKw5qz_0,2863
tokamax/_src/ops/flex_attention/wrapper.py,sha256=RToxUFC-uEXllnIO3rVf_4MYS97E3LWdFAHU6Cik8fQ,3492
tokamax/_src/ops/flex_attention/wrapper_test_base.py,sha256=zWloIOWMe4ZUvmbJMqJyJ5djAu3-mcwk_p1EaDCuiYU,1536
tokamax/_src/ops/gated_linear_unit/api.py,sha256=dzUA-mcmr_NsEUfljuyirRqYOc4R1UsjptAnNT-s0t0,4091
tokamax/_src/ops/gated_linear_unit/api_test.py,sha256=oS4bsKvRxnl-_D4TXd_Wkowafjfu3--bVgLfXfkK3k8,3583
tokamax/_src/ops/gated_linear_unit/autotune_arg_specs.py,sha256=XEUVNtsXI_ykm5wW2QX833NxmY55o8uIupFeS8BfJXM,831
tokamax/_src/ops/gated_linear_unit/base.py,sha256=G00EJ7RJ_ShhHNooYNpqrEYZPpw2AdaIX52Z1myGZWs,5456
tokamax/_src/ops/gated_linear_unit/base_test.py,sha256=zNoOZbucSKmx2up3WFSuuiNG3Nw8wBzs3ekHfm21vUw,2701
tokamax/_src/ops/gated_linear_unit/pallas_triton.py,sha256=lO3S4qXTZU8KkuyR7YJQ3h5LzG3SS-6HpaJ0iQ7e9ok,6949
tokamax/_src/ops/gated_linear_unit/pallas_triton_test.py,sha256=B-btI7moTRg45rMJQ5N3PIFWSnIIC8_hFd9nPN-t7FU,1221
tokamax/_src/ops/gated_linear_unit/test_base.py,sha256=6umFB80b2SLHR2DWTF63qw65UG0W6pxgVylzYl6Srx0,3446
tokamax/_src/ops/normalization/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/normalization/api.py,sha256=XOM18v5RaC9FGPFHQQG2v1FQ1FZLMEZtseS3LOZkbT8,4405
tokamax/_src/ops/normalization/api_test.py,sha256=x9CYz2UkKUlZF75MQ6-LCsfEnxEC3GdWHoojZKWhxmk,5393
tokamax/_src/ops/normalization/autotune_arg_specs.py,sha256=XEUVNtsXI_ykm5wW2QX833NxmY55o8uIupFeS8BfJXM,831
tokamax/_src/ops/normalization/base.py,sha256=a0kS7IwqXufY9g3kTTox25X8JQYhM10PGgY18xPzd54,6171
tokamax/_src/ops/normalization/base_test.py,sha256=XKZTDAMwsol5oEzzdB9Wu-pu2_12eSSadGLyMeKKweg,1027
tokamax/_src/ops/normalization/bench.py,sha256=0UUkQ3n49RgH63iFQ2WZYSKH37DitYXsYY1NaRg95d0,2211
tokamax/_src/ops/normalization/bench_arg_specs.py,sha256=h4sNvnn-BRqDPT9vDJhqau0arnEQYD0UPeEkux_2bf4,2756
tokamax/_src/ops/normalization/pallas_triton.py,sha256=-_WZWvRxlbgTb0VNJJzKKXiGyNM2aHQkddfbFDXDg2A,7762
tokamax/_src/ops/normalization/pallas_triton_config.py,sha256=UA6BR9KnO71kComV5yGEkbysxlQkywJ2ypFqkOoREps,4104
tokamax/_src/ops/normalization/pallas_triton_test.py,sha256=bMG0Zo5oraYn9M-wcVo7Uxp4chb437kFJhvRFZiLNwk,5323
tokamax/_src/ops/normalization/pallas_triton_vjp.py,sha256=G0oSakoV_8B0LZnITzzP-uLK6McIW3DOk5uyZbFZMEY,7052
tokamax/_src/ops/normalization/pallas_triton_vjp_config.py,sha256=CcR6o3tWAXxQYsQyeYJo-x1LLaOwlHOpn-tMqL_4R6s,3092
tokamax/_src/ops/normalization/test_base.py,sha256=tIbEDiLlBTsmDjqvWAxwb1ITgGGgdMl8JPkp4dOJl_0,6662
tokamax/_src/ops/ragged_dot/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/ragged_dot/api.py,sha256=t7g_QToQfEx6Pjf5wymoVlRlfuyV_f2PPNUR1CRbh-4,5042
tokamax/_src/ops/ragged_dot/api_test.py,sha256=M7U1U_sxYLegF_B-EIRUyFPTLAxDh1f_8vT-Lph0bNo,7194
tokamax/_src/ops/ragged_dot/autotune_arg_specs.py,sha256=IQcGwfGxyvSsztMf-KK-BJcTcP2Th2KVYt-kZbNAqXg,910
tokamax/_src/ops/ragged_dot/base.py,sha256=kh0sUMcGAtt6yjGSVyjBTmz9m6CLDwjb1PIAEak-mtM,8905
tokamax/_src/ops/ragged_dot/base_test.py,sha256=vZdsENBA7bp9CKBBr_keoWl0lGVKVWaapqGVzk13HV0,1527
tokamax/_src/ops/ragged_dot/bench.py,sha256=Mb0TWrqeHCIndQDimFxmp2bTwW0kmW2UAHoriMy1gxs,4323
tokamax/_src/ops/ragged_dot/bench_arg_specs.py,sha256=Z_u5bYRvWRCoyJqnINpg1La8S5cuGEugQ699egzfHXo,2033
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu.py,sha256=IFZ_YyGLOAroDRPD_X0O3f7XiOQz102rcu-WeItJYHE,7467
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_common.py,sha256=B8uWkTx_Ef0-IPZFUZ-iKxntyZJ6zLAdvl7LwKGxTKs,9797
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_non_quant_kernel.py,sha256=Ngu1mdzVix9LZlHVWyVccMVPxxq-lT760FXAaRGYsiA,5251
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_kernel.py,sha256=MiXo-k3_zdRlsc79v-z-Pe6hcg60ayK6Asj5ojGdUnU,5057
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_kernel_blackwell.py,sha256=AIa5vsvUQh39jVQleNqLfyKLU-2-LsOfak0nkOCfqWc,14439
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_ws_kernel.py,sha256=gEd4xHCklhEnqy7sEjrHDGFrqYJ1fbFUB-cLPZIcNAs,5877
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_test.py,sha256=ict0JEFXkWoCjIuo_eMSPOCp4MEi2Xv47xcyRkgtfFI,2915
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu.py,sha256=_dMxoaXPZhbXWasz6nj1O_c7G8IGSgUGrprdPaWSz3k,11141
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_kernel.py,sha256=q8xCKq0W4fMXxWhUY3pPw6OKyIKN-tJt09_Gfh4Q84o,30271
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_test.py,sha256=oa9Go3E3_YVBLHkrJ7gEvxNuUX-yyjidXTw1EpjdOeA,4335
tokamax/_src/ops/ragged_dot/pallas_triton.py,sha256=rI_j5ew6Z8_I3KlTcmWVO_BbiE_vAPgNZNTal3I2PE4,14347
tokamax/_src/ops/ragged_dot/pallas_triton_test.py,sha256=9mjlWtehkiU9RG1ULh8CZHAeCMlxQWjAUAeXYEQe06c,2461
tokamax/_src/ops/ragged_dot/test_base.py,sha256=6E1t5Yj6Udw271lW9WYpwxblvs4zyuPG9WIowXY2XFk,7503
tokamax/_src/pallas/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/pallas/block.py,sha256=1DkhMHZvaaZSb7yg1ieR26c_XvOnpOiGFtsCNBuEA6g,11325
tokamax/_src/pallas/block_test.py,sha256=8rrg3tzipZ8Hjp6lVTCRcqUkZF060_3w-ZKp7vftgdc,5462
tokamax/_src/pallas/grid.py,sha256=d_T5YafLTkixhDr8otTNh5frmIcpuHosYrWTbCnqAKQ,3641
tokamax-0.0.2.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
tokamax-0.0.2.dist-info/METADATA,sha256=G2tabkJsxKyPHZ_W5KTvOhkae2X27-vSgiO2neTF8oo,20708
tokamax-0.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
tokamax-0.0.2.dist-info/top_level.txt,sha256=xHgO7qOuYjrLH3047bLQLnwrjI76i9-xlChxNQIZn-Q,8
tokamax-0.0.2.dist-info/RECORD,,
