Index

A

Accelerator (class in xma.accelerator) autotune() (in module xma.autotuner.tuner)

AutotuneConfig (class in xma.autotuner.config) AutotunedFunction (class in xma.autotuner.tuner)

B

backward() (CustomOp static method)

bmm() (in module xma.functional.bmm)

C

ceil_divide() (in module xma.math) check_power_of_2() (in module xma.math) continuous_count() (in module xma.functional.continuous_count) cpp_jit() (in module xma.jit) cpu (Accelerator attribute)

cross_entropy() (in module xma.functional.cross_entropy) ctx_needs_gradients() (in module xma.custom_op) ctx_save_for_backward() (in module xma.custom_op) cuda (Accelerator attribute) (KernelBackend attribute) CustomOp (class in xma.custom_op)

D

divide_if_divisible() (in module xma.math)

down_projection_triton_forward() (Experts method)

E

enable_counters() (in module xma.counters) enable_kernels() (in module xma.inductor) Experts (class in xma.layers.moe)

exposed_signature (AutotunedFunction property) extra_repr() (Experts method) (GRU method) (RNN method)

F

forward() (CustomOp static method) (GRU method) (LinearAttention method) (M2RNN method) (MoE method) (RNN method)

forward_backward_torch() (CustomOp static method) fused_linear_cross_entropy() (in module xma.functional.fused_linear_cross_entropy) fused_residual_add_rmsnorm() (in module xma.functional.fused_residual_add_rmsnorm) fused_residual_add_rmsnorm_backward() (in module xma.functional.fused_residual_add_rmsnorm) fused_residual_add_rmsnorm_forward() (in module xma.functional.fused_residual_add_rmsnorm)

G

get_accelerator() (Accelerator static method) get_cartesian_product_autotune_configs() (in module xma.autotuner.config) get_compatible_accelerator() (KernelBackend method) get_core_count() (Accelerator static method) get_counter_value() (in module xma.counters) get_current_device() (Accelerator static method) get_fused_residual_add_rmsnorm_replacer() (in module xma.inductor)

get_kernel_backend() (Accelerator static method) get_key_values() (AutotuneConfig method) get_next_power_of_2() (in module xma.math) get_powers_of_2() (in module xma.math) get_rmsnorm_replacer() (in module xma.inductor) GRU (class in xma.layers.gru.module) gru() (in module xma.layers.gru.op)

I

increment_counter() (in module xma.counters)

is_condition_valid() (AutotuneConfig method) is_counter_enabled() (in module xma.counters)

K

KernelBackend (class in xma.accelerator)

L

linear_attention() (in module xma.layers.linear_attention.op)

LinearAttention (class in xma.layers.linear_attention.module)

M

M2RNN (class in xma.layers.m2rnn.module) m2rnn() (in module xma.layers.m2rnn.op) module xma xma.accelerator xma.autotuner xma.autotuner.config xma.autotuner.tuner xma.counters xma.custom_op xma.functional xma.functional.bmm 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.inductor xma.jit xma.layers xma.layers.gru xma.layers.gru.module xma.layers.gru.op xma.layers.linear_attention xma.layers.linear_attention.module xma.layers.linear_attention.op xma.layers.m2rnn xma.layers.m2rnn.module xma.layers.m2rnn.op xma.layers.moe xma.layers.rnn xma.layers.rnn.module xma.layers.rnn.op xma.math xma.optimizers xma.optimizers.sgd xma.optimizers.sgd.module xma.optimizers.sgd.op

MoE (class in xma.layers.moe) mps (Accelerator attribute) (KernelBackend attribute)

N

nki (KernelBackend attribute)

P

p_norm() (in module xma.functional.p_norm) pack_sequence() (in module xma.functional.sequence_packing)

pallas (KernelBackend attribute) partialize_and_update_signature() (in module xma.inductor)

R

reset_counters() (in module xma.counters) reset_parameters() (Experts method) (GRU method) (M2RNN method) (RNN method)

rmsnorm() (in module xma.functional.rmsnorm) RNN (class in xma.layers.rnn.module) rnn() (in module xma.layers.rnn.op) rocm (Accelerator attribute) (KernelBackend attribute) run() (CustomOp class method)

S

SGD (class in xma.optimizers.sgd.module) sgd() (in module xma.optimizers.sgd.op) softmax() (in module xma.functional.softmax)

step() (SGD method) swiglu() (in module xma.functional.swiglu) swiglu_packed() (in module xma.functional.swiglu) synchronize() (Accelerator static method)

T

torch (KernelBackend attribute) torch_forward() (Experts method)

tpu (Accelerator attribute) trainium (Accelerator attribute) triton (KernelBackend attribute)

U

unpack_sequence() (in module xma.functional.sequence_packing)

up_projection_triton_forward() (Experts method)

V

verify_accelerator() (KernelBackend method)

X

xma module xma.accelerator module xma.autotuner module xma.autotuner.config module xma.autotuner.tuner module xma.counters module xma.custom_op module xma.functional module xma.functional.bmm module xma.functional.continuous_count module xma.functional.cross_entropy module xma.functional.fused_linear_cross_entropy module xma.functional.fused_residual_add_rmsnorm module xma.functional.p_norm module xma.functional.rmsnorm module xma.functional.sequence_packing module xma.functional.softmax module xma.functional.swiglu module xma.inductor module xma.jit module

xma.layers module xma.layers.gru module xma.layers.gru.module module xma.layers.gru.op module xma.layers.linear_attention module xma.layers.linear_attention.module module xma.layers.linear_attention.op module xma.layers.m2rnn module xma.layers.m2rnn.module module xma.layers.m2rnn.op module xma.layers.moe module xma.layers.rnn module xma.layers.rnn.module module xma.layers.rnn.op module xma.math module xma.optimizers module xma.optimizers.sgd module xma.optimizers.sgd.module module xma.optimizers.sgd.op module xma_op() (in module xma.custom_op)