tokamax/__init__.py,sha256=R7F2RUV68xfnVXfz0j4PV9YUR6aHj_X7DFddkVxvHhA,1550
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=lP8pe0GH7_ST7BVvVgacqxekKlER2zoR9BcuM86RfIA,3966
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=FokHPoMVQrr7uhD1SrtECitKaYf9BkfQ_YcNBSpVK5s,7913
tokamax/_src/batching_test.py,sha256=voxWU-AjoDFXIAoSeI6HA9peNntiP_2jueETvKh959Y,4127
tokamax/_src/benchmarking.py,sha256=kj6wUl_xFbs6rLZQbkk4cgaX0sh3m5H5C4mgfVP0vMo,17478
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=dMdUp9iCOrWuIm7TXry95gYR2yB2TmNDC8VTcqdlNiY,1051
tokamax/_src/hlo_utils.py,sha256=0xxBUVsbydCj23h9qevpDNPtuyL7XXDRhi-hro0dpLk,12995
tokamax/_src/hlo_utils_test.py,sha256=xY5SlM7Ayz-Ob9x1t8CXhfqe83uLWONp8fYZdgDwJcI,17312
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=NhEtMdrOeKHxW8Jxj9MNVVlXBf_3pH5jOpmpauAj3mk,2364
tokamax/_src/numerics.py,sha256=QN7yE1rlQt1novbEYsUTSR_9wyrlIm4BHf0pnjDHnwc,6314
tokamax/_src/numerics_test.py,sha256=GpW5WfhS3YE435Np_UkW37R1jCc0R-hlHJrRaePIvIk,4884
tokamax/_src/precision.py,sha256=Oy1kicxjgmS3TOl9lRwXhQ0C7oyXPZTLBXwE0nuSEGE,5722
tokamax/_src/precision_test.py,sha256=Cpba1dPi8mKK22gfb4fCyAd1NE62MAd11Du3GUx5DRc,3584
tokamax/_src/pydantic.py,sha256=DwX7yrOp17VCcWMhGPDQlMFwODoR7T16U8YTuUKnY60,10168
tokamax/_src/pydantic_test.py,sha256=qLUDtvMDHT_YpegsBnQjMCsSu5ZkxwYyTpLYnHF47ww,7727
tokamax/_src/quantization.py,sha256=UuOvAOSawTgmR9gsZnKLSe687r8UYILumuMEXqM3JnU,4090
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=WLVFYIumiULoeWGWrFA_pNKXaI1MrF0P6WoMMpgfJL4,10397
tokamax/_src/autotuning/api_test.py,sha256=dcEUSYneLTiavFZxWbj4REJ5hnVnsd1Duwelb5mQpUM,12808
tokamax/_src/autotuning/autotuner.py,sha256=zZ69C9YU0edYwTgHptIJdgYuOR61ioPCPyC0OCjRYhc,6732
tokamax/_src/autotuning/cache.py,sha256=2m5tO3IIOIq1CjJ2qp9qZrE-1yJwCV4V29SbHbB3EY8,1553
tokamax/_src/ops/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/op.py,sha256=u_OVYj6bI9jiKd22UJwTMXnpOELRCjJWRVZi4v3fWQY,20776
tokamax/_src/ops/op_test.py,sha256=zlkyZLzesTFwN5f9u12X5ASCKXN4o-YKDK7rDFHu8uc,7587
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=tAlmMpy3blgIyO4wbEP24eBj2N2Yck7ma1g9IE9MxWA,13643
tokamax/_src/ops/attention/base.py,sha256=RQvlBWM5l4jQrnWxR3_sPl2KjxqT5qtjdQjQFUPHLu8,28809
tokamax/_src/ops/attention/base_test.py,sha256=L9gAlnSReP0oyNL-lpQ7i9yrXcyfMoPV2NilwdCYjCo,5286
tokamax/_src/ops/attention/bench.py,sha256=_wACi8pCfJo5g8TtHXfYIEVjLvRygOJrr-HMWuFOHWY,2606
tokamax/_src/ops/attention/bench_arg_specs.py,sha256=fC0RJ4sPYnCQB_VX-dP95vfIz6rxjARWwcgDHjhaLiI,1437
tokamax/_src/ops/attention/jax_nn.py,sha256=qUVr_C9habnOvfXanYVx1U2R0JmcWmONhZ-_Nt-GTuo,5532
tokamax/_src/ops/attention/jax_nn_test.py,sha256=f4imkbLrA_UIaQ-8QABdvmGP3-D0D3bLcdp-JlDnf6w,3927
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=rVSihtvZj_ak3WdQXWtFJkH738WboNbLJMGHSm6ZOC8,4681
tokamax/_src/ops/attention/pallas_mosaic_gpu_flash_attention_vjp.py,sha256=qMvGb6T-jzDoo5eSdRjiMJ7bo1eCPV6HAeXaFAbKJBs,27529
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=clLLTDu799rlGj95N7hVAmfq5o0W-AhrDjZ7JZizyaY,36099
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/flex_attention/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/flex_attention/base.py,sha256=02HBWZLyEbBjrinOoDINDtHyYELULmVFPE71P4ZDR6Y,12222
tokamax/_src/ops/flex_attention/base_test.py,sha256=Ea9iviHpoPJ6spG7ZquCp6VV2H2qOTNYLjo_WSYSapI,1265
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=owaQ9D8TqbufhOShKnPdX0QY3QDpWS6C453VUOSTrMI,3697
tokamax/_src/ops/gated_linear_unit/base.py,sha256=PGz43efj5F-oEp0sIMNOWpPWu7BSqtBz15PXj5NupTg,4968
tokamax/_src/ops/gated_linear_unit/base_test.py,sha256=zNoOZbucSKmx2up3WFSuuiNG3Nw8wBzs3ekHfm21vUw,2701
tokamax/_src/ops/gated_linear_unit/pallas_triton.py,sha256=2P8Bd8m-PwBegUrfpL7M1lAGyzEDxmQEzPbjQ-KcIlA,6960
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/base.py,sha256=a0kS7IwqXufY9g3kTTox25X8JQYhM10PGgY18xPzd54,6171
tokamax/_src/ops/normalization/base_test.py,sha256=XKZTDAMwsol5oEzzdB9Wu-pu2_12eSSadGLyMeKKweg,1027
tokamax/_src/ops/normalization/bench.py,sha256=bRudd6UuTQc_qcYZbRGP9lDoyPEsgkHZ_VR-wOFERdg,2165
tokamax/_src/ops/normalization/bench_arg_specs.py,sha256=YSWRz1iZhtdfMtAOsA--Lv1BJkCNZhspTqxQTzd4hJc,2115
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=xvtIwdRI7or2sLM-TTj7noQSpWNSCrQglbj75kFN8Kw,6595
tokamax/_src/ops/ragged_dot/__init__.py,sha256=sojZirG99P27D5bJ9BjPUNR3NbS1SYi4vZMSjOBmGHM,696
tokamax/_src/ops/ragged_dot/api.py,sha256=QITkayTb8ANIdh8SKi77CWhF2rrW5SnRNb60ZQ_DWc8,4427
tokamax/_src/ops/ragged_dot/api_test.py,sha256=YvzZ8QgsMgXGcYGHqbJkU2_An591nQTYMvHi0TSkP6w,7256
tokamax/_src/ops/ragged_dot/base.py,sha256=BSfBH51ovX2qMeRZK31Yc3yDzpI2U1a5xKOnKofICvs,8542
tokamax/_src/ops/ragged_dot/base_test.py,sha256=vZdsENBA7bp9CKBBr_keoWl0lGVKVWaapqGVzk13HV0,1527
tokamax/_src/ops/ragged_dot/bench.py,sha256=8y-abSGib_jYDMPHs3VDe65d9EZJ5nfFaxLDCO7N2EA,4278
tokamax/_src/ops/ragged_dot/bench_arg_specs.py,sha256=S2EjpoVirG1njR9ODHmeNzhDfdFPTWRnRlRwZFQOrgo,1796
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu.py,sha256=LdcvBbmjdw4IVvznaYeIVsGTBaS6Yu5zhHGCbay01cU,7246
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_common.py,sha256=k1A4Tf034DYeuJqcGYDMmRYeoPwSY-guswQw3ie_aZ0,9945
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_non_quant_kernel.py,sha256=DE6StEOm1bC1wRfFFbHvWBtOX0EDD6lS45kaSJ3H5JQ,5252
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_kernel.py,sha256=1zI8UmRZLola_LrMAF5gz-egrPLXmbPWcAYAPx4CPX0,5065
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_kernel_blackwell.py,sha256=mq7U8iCSnrm88y6DFTPGTNjwpHxWPKf_edRCzNVDiXQ,14463
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_quant_ws_kernel.py,sha256=JR4ABAxt4tppwBVmnhHohlgSU9OksRkRIlQHzfCFmAc,5885
tokamax/_src/ops/ragged_dot/pallas_mosaic_gpu_test.py,sha256=FZuSQYpRwhaegrI6mkiSvJwvSDZsbFaXJyiPuebGpxo,3662
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu.py,sha256=O-SZOPt5ulDcuUxOftlsyOGo0OguCO38ObJnWZvuCns,5409
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_kernel.py,sha256=_NJU62YaCHmCwbGKuHsmhJl3HSWfjOw8PnoCEFWQRYw,27651
tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_test.py,sha256=HBeOegWHerYWQGCAy-w8vuJ8VTh7_RMTRA1TyvONJr4,4804
tokamax/_src/ops/ragged_dot/pallas_triton.py,sha256=9ZWvI2Ncubej1vlTkPE-xpzWYsqRQauhUaCHl6nV1xE,14356
tokamax/_src/ops/ragged_dot/pallas_triton_test.py,sha256=Nuv47I2jDCophJE6YeAk1vlMo8R33Kx0syJHV84OcAA,2517
tokamax/_src/ops/ragged_dot/test_base.py,sha256=jWJkyYT-NxhQJ9mXkVZ1TQi78KG6wMybDpogAH2nPV0,6672
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.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
tokamax-0.0.0.dist-info/METADATA,sha256=EaC3oBhfYLb8PeQaplX9DJ8vBn4q9mjSVFz1u6xrYKs,20685
tokamax-0.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
tokamax-0.0.0.dist-info/top_level.txt,sha256=xHgO7qOuYjrLH3047bLQLnwrjI76i9-xlChxNQIZn-Q,8
tokamax-0.0.0.dist-info/RECORD,,
