tokamax/__init__.py,sha256=kYCnTL6FBb3u8O28AT0b51kmPYgne5dULDCHUg_xNHg,1877
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=KIvhgep9hty_sGz6MwA2WYjOPr622Wym5kYfrA8BVX8,4268
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=xohJP_n17BNFzmCIXlXBFWZwUUS83r2t-lG9ivdVtVk,8007
tokamax/_src/batching_test.py,sha256=M3WGauph8yoMEI7agmRi2W-BWkA_osnXntftArDXRRI,4493
tokamax/_src/benchmarking.py,sha256=TDSTmhsqJ8cIxBSgAEDRQqkPlmxdAQGB0jw07O9O3z0,18026
tokamax/_src/benchmarking_test.py,sha256=GhVUYffmJAoZc0lvvUUfHkAr1dGZoHeW1ehvUtAtHzg,8312
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=98WSJRzAP-iSrYTrfw-R8d1GVKJup4aUSwnriXIcuko,1147
tokamax/_src/mosaic_tpu.py,sha256=5-eCQ0EPe8KfenIYcBcUtTF5cl9qgx1Q44ncRM0o8DM,7847
tokamax/_src/numerics.py,sha256=k-srlztc3Ea8AhUGCct_A6ln-BJ9TiSJDdBDHnxeapc,7002
tokamax/_src/numerics_test.py,sha256=860LLuVtjgF7zOswGir56VUjEK-FhYuYhv_K7UFeJKA,6329
tokamax/_src/precision.py,sha256=BVYnC6VtC3wk4arNUcTujf4LtOKpGf4dQe7wxIl9gMM,7209
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=hE5b3Kh9rAluThjCgjjr12cYXO1RUSlYH-r9xlX3B6Q,4710
tokamax/_src/quantization_test.py,sha256=cFEJrvZjat4TEjz2kERRxdA4J4C3fRVqtVbwbmn83TA,3133
tokamax/_src/shape.py,sha256=772lVlGbXfDgKiUrTSP_jlYKSPn0HSOsVIVwCExK0bU,2935
tokamax/_src/shape_test.py,sha256=8C6awpma4oJzPlcIpqwVvf5SSfKmATfghDA6hZkr5Os,1560
tokamax/_src/test_utils.py,sha256=QKNNgaGOBMNbaO-IUL1gbXydachSMhmIW2URPiRtoSE,2023
tokamax/_src/triton.py,sha256=mp-Hzqn07R55vuNe3JdfWikfE7VGiRGtnNGTeguXXMo,1428
tokamax/_src/triton_test.py,sha256=0Y-4nft57BwZvII3JZJTu6BjdCu893zj7D9hrcVAmhM,1045
tokamax/_src/utils.py,sha256=33UeGzaHrnfuMmL6DHqphcecRXp1RJUbkFHENhx1kYQ,1944
tokamax/_src/autotuning/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/autotuning/api.py,sha256=Zrl_q51UZ6zeN7nJxGhukY21ZqLQcPHuHK8bjC8OnNA,10279
tokamax/_src/autotuning/api_test.py,sha256=MS2_IMHKO_bh6JvrQhmy9hE0FGJUSHkLhHOLYtoDRsY,12857
tokamax/_src/autotuning/arg_spec.py,sha256=MQiziGpPRlXMVahqxQwtQ5pyCfN4aepIcthctgs6LUQ,1870
tokamax/_src/autotuning/autotuner.py,sha256=lyBfvUKF5xTQ4kkFgq6A-0tHrn-i0wexz49gnYwJC5Y,6752
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=r1_1Yeqb2Fvm6HdIt7WD5keuXrD4J-pVbAy3SyjU5SQ,22830
tokamax/_src/ops/op_test.py,sha256=BaQeYvnWutT51p94tiHpfR8UuBIm2UoL_GWhxE3yJSs,7706
tokamax/_src/ops/attention/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/attention/api.py,sha256=H2oU39CbgGRyTvzFsh1ufrWS4bk-Oy6ft7C4iB33wZ4,6202
tokamax/_src/ops/attention/api_test.py,sha256=f5hsdxa2ZyPQoEViGvD_nJPIdifO4Nsawi9bv6Hm4Cg,11785
tokamax/_src/ops/attention/autotune_arg_specs.py,sha256=uCEz3-up3KLU49P9Tny34QeYecA9DZ39MLEOygsvp_Y,3438
tokamax/_src/ops/attention/base.py,sha256=PQ1FbSi97fUV_kg67YmEIBGcNquMXWMgzF-jZ4Y6LE4,28821
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=6mwgyWMTLroYcEJqLE8ljikoaDqpq8F471nGeXWMiPw,1687
tokamax/_src/ops/attention/jax_nn.py,sha256=2pdemyrP9eKeC7M7rTObptyssTk82G4z1EjUvTDpnWg,5775
tokamax/_src/ops/attention/jax_nn_test.py,sha256=WNRmKjHjYrRIFZfOm74c849My7pAqPeKvGdaiqrI0-8,3809
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention.py,sha256=pMlneoqtdtbYH-vACPOK12QC9YblmfVGLmBVBDZHnlM,25275
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention_test.py,sha256=LmFA0r1RDkctkB0kOwyJm_MbJHttJqRRFHOo2nQ4XFQ,5990
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention_vjp.py,sha256=FtZbTxNL3PwxZeAs0Pnsn2pISkX2Y_nDFfMLJuDWSk4,30794
tokamax/_src/ops/attention/pallas_triton_flash_attention.py,sha256=EeAIRYHgiiLXRBPI_z6gmZP_Dj0sVNawiAub_QGxcKo,24050
tokamax/_src/ops/attention/pallas_triton_flash_attention_test.py,sha256=7cDzk64yAMFKwIn3JUQZ4LivNcOf__9ovu1izG04gyI,5793
tokamax/_src/ops/attention/pallas_triton_flash_attention_vjp.py,sha256=rrNXkuthqvAO7u9lNj5TCWwiNX9gtMIuM5DHaxv-Tkg,18076
tokamax/_src/ops/attention/test_base.py,sha256=EupLD9yh8uJ5rjL_t78Z47agY_kUuXHS4JiZ6K_7yNQ,36187
tokamax/_src/ops/attention/xla_chunked.py,sha256=p36mtHi_EJM2M8D0sIUH_VZXd9HPaZFzPTDGNrcV2g4,12352
tokamax/_src/ops/attention/xla_chunked_test.py,sha256=n6Eg9KAq3Jv0BZS89G4bComLdWzhunISXbUQw49UJbs,9078
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel.py,sha256=rPR3d57JjbmNbHidF6oUcpR_BfG2VafvEly9KcsmAl4,74570
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=d7R6mva-EPZxLVgIoeHXqICfDHDeZI6KXu5wV1gS2dw,22371
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=P2E3eoMRK-w8hhjN-swzwL-LicdJ4S4XKVWYz_iJjKk,20148
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask_test.py,sha256=cla9mve6li7SUu0sC_-dsdHFMGnecZepLKtpwPniVvU,46500
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=3ktlcBfBtuvigzRz9M11Jjfqnem0iZDoTIH686FoFcw,815
tokamax/_src/ops/flex_attention/base.py,sha256=JsttyYD9zW5Ph6bwaZp9ZMwq47-47NXqICa5q-EYRqg,12351
tokamax/_src/ops/flex_attention/base_test.py,sha256=l4y8geGYw4AIJ_HnNOfXIyHXw2LKVc1xYEBlwNhIws0,1475
tokamax/_src/ops/flex_attention/pallas_triton.py,sha256=EmaYAbrT9y6LVEQqqolMpOz898xOujMqG3CLg5Ej3Sk,18436
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=q54LZB_rEPukcHyypWDpntwKJT-7mEMspOvRKQe81YY,3543
tokamax/_src/ops/flex_attention/wrapper_test_base.py,sha256=fRjIzpbm6y1HXTgHk6JiDrGRo5wSOqk45U12DLQyo5c,1535
tokamax/_src/ops/gated_linear_unit/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/gated_linear_unit/api.py,sha256=P_P1bmK-kK6b4Dsc0giw4b9rpVnjjZkbuPxKWe467FY,4278
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=Y70pdt99ZhbqcGChBTpDDINynbntdJmRyIGfm6IU3_w,823
tokamax/_src/ops/gated_linear_unit/base.py,sha256=Pd7PO1hQbrYex1NiCLjB4dVHUQ-AJXS965JNamge4GI,5531
tokamax/_src/ops/gated_linear_unit/base_test.py,sha256=zNoOZbucSKmx2up3WFSuuiNG3Nw8wBzs3ekHfm21vUw,2701
tokamax/_src/ops/gated_linear_unit/pallas_triton.py,sha256=FWOd0SHb6-lUoLl_XtBYb1pAvFgX6NRXqXSRYppQ19I,7175
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=JLjrk-4xOiep3uYHT93bq2bN6o4eMUMUKn7AkRY_1Qo,4603
tokamax/_src/ops/normalization/api_test.py,sha256=zgqj2IGx9X82sFGxVTIzWmDNTbhpR1O6uMlYXNzimwg,5359
tokamax/_src/ops/normalization/autotune_arg_specs.py,sha256=Y70pdt99ZhbqcGChBTpDDINynbntdJmRyIGfm6IU3_w,823
tokamax/_src/ops/normalization/base.py,sha256=Wyd19A7RNgC3fcHNBli_y1jFt2Y2PVnYyk9vXmccsb4,6246
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=trQYL9acYBYVBO16mdozgBeHbWSUSlDuKljK6UZ4nAs,2739
tokamax/_src/ops/normalization/pallas_triton.py,sha256=lAfUkf_UwlK3Q8sxfWuQ_Dr8oI1c8aOMKSCKseMDFSw,8011
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=nn-qdJGkPiMRW9vcJgGH5SFJPW7thNgpit3hlwh3l4Q,7301
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=q_WIXRi7m1b8UNxSoA1B7mrFEFUpMh4S_KXFFmxVOow,902
tokamax/_src/ops/ragged_dot/base.py,sha256=C1v1haPxe6gtZ0462mPSWh6BM8mD6jNrecjJikPvwBk,9714
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=flTIGBPWvN7fkXGWJ-Zr1mMIGu6qa4bcCNCjvqqxbww,2017
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu.py,sha256=FFXq6KOjEde3OHAO4o9xwHnlcRa8qi1mGfgEo6ER42U,7666
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_common.py,sha256=bGswRLr1FRBj2etKoIfs0Tl1mhkAgfsBwaWUKivRvNs,9843
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=AlJkbgtJNfHLLAV9pJniRT0zncCn7GmlsZFJvZq4Xok,4809
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_kernel_blackwell.py,sha256=OUyT1V8CNwTkZMUwbb2us-WEJLPPnQJ7NFsxxpgHSuQ,14338
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_ws_kernel.py,sha256=9Y7pF1-dml1AFPaxrjVnL2yUREXFpdj3agwSZyIxh-k,5619
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_test.py,sha256=ict0JEFXkWoCjIuo_eMSPOCp4MEi2Xv47xcyRkgtfFI,2915
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu.py,sha256=V-5a6y-O60zA4eBSwdqT_jX4eaUyQ8ouoqjX2Axsi8U,12743
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_kernel.py,sha256=EGaScL2s9bTMULSx7UfwD4qMIxyb1906YgTw3_TFAB8,25047
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_test.py,sha256=UiBLYg_X2lNwqTEIxdJu7pgtrX8mXnv2A06NE3Qz9L4,4729
tokamax/_src/ops/ragged_dot/pallas_triton.py,sha256=tcPsSYGKRMUr00D7UxRq0lHq_Can0HJ4pVts0O_Fqbc,14539
tokamax/_src/ops/ragged_dot/pallas_triton_test.py,sha256=9mjlWtehkiU9RG1ULh8CZHAeCMlxQWjAUAeXYEQe06c,2461
tokamax/_src/ops/ragged_dot/test_base.py,sha256=nqUx4hVtKrSi8lPKEaq5h6tLiClJ8mwZgjtXIM3suqU,7727
tokamax/_src/ops/triangle_multiplication/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/triangle_multiplication/api.py,sha256=6JPpNLmwOTTk5ML4q385Xay0TEnrrwzWtPQiarJoPcs,5060
tokamax/_src/ops/triangle_multiplication/api_test.py,sha256=21FGqXlqLUCVW4bvMFKewI7ojWRMMVXNSWu_FHgeFmQ,2831
tokamax/_src/ops/triangle_multiplication/autotune_arg_specs.py,sha256=3ktlcBfBtuvigzRz9M11Jjfqnem0iZDoTIH686FoFcw,815
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.4.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
tokamax-0.0.4.dist-info/METADATA,sha256=6qEQibhtPEBtReadfV9G8SjJkFs9vSUD4HQMIvUXN3c,20981
tokamax-0.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
tokamax-0.0.4.dist-info/top_level.txt,sha256=xHgO7qOuYjrLH3047bLQLnwrjI76i9-xlChxNQIZn-Q,8
tokamax-0.0.4.dist-info/RECORD,,
