{"library":"tokamax","type":"library","category":null,"description":"Tokamax is an OpenXLA library providing high-performance custom accelerator kernels for NVIDIA GPUs and Google TPUs. It offers state-of-the-art implementations built on top of JAX and Pallas, along with tooling for users to build and autotune their own custom kernels. As of version 0.0.12, it is still under heavy development, and users should anticipate API changes.","language":"python","status":"active","version":"0.0.12","tags":["JAX","GPU","TPU","Custom Kernels","Deep Learning","Performance","Pallas","OpenXLA"],"last_verified":"Mon May 25","install":[{"cmd":"pip install -U tokamax","imports":["import tokamax","import jax","import jax.numpy as jnp","tokamax.layer_norm","tokamax.dot_product_attention","tokamax.autotune","tokamax.standardize_function","tokamax.benchmark"]},{"cmd":"pip install git+https://github.com/openxla/tokamax.git","imports":[]}],"homepage":null,"github":null,"docs":null,"changelog":null,"pypi":"https://pypi.org/project/tokamax/","npm":null,"openapi_spec":null,"status_page":null,"smithery":null,"compatibility":{"summary":{"python_range":"3.10–3.9","success_rate":15,"avg_install_s":19.4,"avg_import_s":11.39,"wheel_type":"wheel"},"url":"https://checklist.day/v1/registry/tokamax/compatibility"}}