Contents Menu Expand Light mode Dark mode Auto light/dark, in light mode Auto light/dark, in dark mode Skip to content
XMA documentation
XMA documentation

API Reference

  • xma.functional
    • xma.functional.continuous_count
    • xma.functional.cross_entropy
    • xma.functional.fused_linear_cross_entropy
    • xma.functional.fused_residual_add_rmsnorm
    • xma.functional.p_norm
    • xma.functional.rmsnorm
    • xma.functional.sequence_packing
    • xma.functional.softmax
    • xma.functional.swiglu
  • xma.layers
    • xma.layers.depthwise_causal_convolution
    • xma.layers.gru
    • xma.layers.linear_attention
    • xma.layers.m2rnn
    • xma.layers.moe
    • xma.layers.rnn
  • 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
Back to top
View this page

xma.layersΒΆ

  • xma.layers.depthwise_causal_convolution
  • xma.layers.gru
  • xma.layers.linear_attention
  • xma.layers.m2rnn
  • xma.layers.moe
  • xma.layers.rnn
Next
xma.layers.depthwise_causal_convolution
Previous
xma.functional.swiglu
Copyright © 2026, Mayank Mishra
Made with Sphinx and @pradyunsg's Furo