tokamax

import jax import jax.numpy as jnp import tokamax def loss_function(x, scale): # Apply layer normalization with a Triton implementation x = tokamax.layer_norm( x, scale=scale, offset=None, implementation="triton" ) # Apply dot product attention, allowing Tokamax to select the best implementation x = tokamax.dot_product_attention(x, x, x, implementation=None) return jnp.sum(x) # Example usage with JAX JIT and Grad key = jax.random.PRNGKey(0) x = jax.random.normal(key, (32, 2048, 64), dtype=jnp.bfloat16) scale = jax.random.normal(key, (64,), dtype=jnp.bfloat16) f_grad = jax.jit(jax.grad(loss_function)) output_grad = f_grad(x, scale) print("Computed gradient successfully.") # Example of autotuning (requires compatible hardware) # autotune_result = tokamax.autotune(loss_function, x, scale) # with autotune_result: # out_autotuned = f_grad(x, scale) # print("Autotuned output successfully.")