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=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=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=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=_Tkn7fojxqQ643p54MVuRTeaBYQtXGxmLnAjTkD340I,5852
tokamax/_src/ops/attention/api_test.py,sha256=f5hsdxa2ZyPQoEViGvD_nJPIdifO4Nsawi9bv6Hm4Cg,11785
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=9ZBtVwh0EDdgiAhoPD3tgiabkI1K2oxC2eCLUdxN2jk,5712
tokamax/_src/ops/attention/jax_nn_test.py,sha256=WNRmKjHjYrRIFZfOm74c849My7pAqPeKvGdaiqrI0-8,3809
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention.py,sha256=EUMOsKdeQ3YIKdWrcl8-F4GSJk1_gopmNSAhwGfhwms,24798
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention_test.py,sha256=UNRfN5khbl57oeVJpjjx3jCaEcaHvOy3BBz6rh2ufeQ,5335
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention_vjp.py,sha256=9lpe_QLkmw3F9H101CcrC1087FVhKkH0_aHo0WIImxs,30087
tokamax/_src/ops/attention/pallas_triton_flash_attention.py,sha256=ayK9t0Aqrjy2O5VaohVJT-1wihFxMA1xOsCxOe_V2YQ,23962
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=YxNkZi45IQMP9KDj4ftss-YffQDcYZZuC59dltfT2yw,17618
tokamax/_src/ops/attention/test_base.py,sha256=5ks7ywD2nb0yBV3OSnAvLcVpotuLmnTg1Afg3JllBEo,36371
tokamax/_src/ops/attention/xla_chunked.py,sha256=WhIvaBH-LDIRiCi6Kzkl8xcZ0La5yHPnY6XdZrJNW40,11483
tokamax/_src/ops/attention/xla_chunked_test.py,sha256=_y92iukPWgv5Yrdu5yE5lLTWKm4ErSMlVT_vKiDyd3A,8984
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel.py,sha256=xFlHXLu1JQoWql4ShTg2PVmq1f5_SiwnMEjIvK99xFU,74267
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=Nri5mJwmOgIkAxe3aVnMqLEsuSHySJLCkMITux8xpvw,20076
tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_mask_test.py,sha256=ruIOY9OA5I2BoorUIhHH80oRpYiqrxtFgH-gkTEoIQY,47096
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=fdQBf0FvnmG5IS9hocPanC7hgqBaMPXS0X0AvzCuyHU,12288
tokamax/_src/ops/flex_attention/base_test.py,sha256=l4y8geGYw4AIJ_HnNOfXIyHXw2LKVc1xYEBlwNhIws0,1475
tokamax/_src/ops/flex_attention/pallas_triton.py,sha256=KagHOBDvTL_7rtRrYFFu2tB-EfAOYZ3-2mG49aix6L0,18348
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=fRjIzpbm6y1HXTgHk6JiDrGRo5wSOqk45U12DLQyo5c,1535
tokamax/_src/ops/gated_linear_unit/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
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=gBVUVskcZDwT-vCbX31eYqSGlhBXyABcP_scx1dinOw,7099
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=C3lATx-bjx9eNXuWeX-ShInYK_Q_rNsGZsxcXQ4xYbw,7912
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=7i7FVkMI3fW6MKzSdqfzRxuEDQSUj2pIGP8nkwToPUw,7202
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=NZYZJTfNGwaoe8NV1QEEwD-MboEWCd3eJnXkkdf1qzE,9651
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=5VANARoK-uwyQRPVI3NyK-FR40wOT0-4aL047rbrTjc,7579
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=TEycT_iupTe-uYQtlFVWWAdmkjJyJRMd2zxgKxoKRdk,4836
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_kernel_blackwell.py,sha256=_sENQNck-CMC92kzNKth4oCWq2r0B2pUt_MR039T8y8,14473
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_ws_kernel.py,sha256=dat_YPtU21HeYcQrw6Io5o4SzQwTiF-VHYh0o9rxM5o,5646
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_test.py,sha256=ict0JEFXkWoCjIuo_eMSPOCp4MEi2Xv47xcyRkgtfFI,2915
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu.py,sha256=2TQh8diD0BxZl4si7Y2jOuCbQ1aqR_zO3hI29dbJhkw,12306
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=4qyEHJ5YzZgIJl6yH1CkRm8cH83_qe3caBiKNnmxvV0,4261
tokamax/_src/ops/ragged_dot/pallas_triton.py,sha256=b_ol5oDBtShr3A2TDQ8r7uk4ue37WzESC84f0Cc-qqw,14451
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=XEUVNtsXI_ykm5wW2QX833NxmY55o8uIupFeS8BfJXM,831
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.3.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
tokamax-0.0.3.dist-info/METADATA,sha256=_rdfEFx54wbKrG9Qtwp6pCAVH7Cs-9kukrf7ijo36So,20941
tokamax-0.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
tokamax-0.0.3.dist-info/top_level.txt,sha256=xHgO7qOuYjrLH3047bLQLnwrjI76i9-xlChxNQIZn-Q,8
tokamax-0.0.3.dist-info/RECORD,,
