Skip to content

XMA documentation

XMA documentation

API Reference

xma.functional
xma.layers
xma.layers_jax
- xma.layers_jax.depthwise_causal_convolution
- xma.layers_jax.linear_attention
xma.optimizers
- xma.optimizers.sgd
  - xma.optimizers.sgd.param_group

Utilities

xma.accelerator
xma.counters
xma.autotuner
- xma.autotuner.config
- xma.autotuner.tuner

xma.functional.softmax¶

xma.functional.swiglu

xma.functional.sequence_packing

Copyright © 2026, Mayank Mishra

Made with Sphinx and @pradyunsg's Furo