xma

• xma.accelerator
  • Accelerator
    • Accelerator.cpu
    • Accelerator.cuda
    • Accelerator.get_accelerator()
    • Accelerator.get_core_count()
    • Accelerator.get_current_device()
    • Accelerator.get_kernel_backend()
    • Accelerator.mps
    • Accelerator.rocm
    • Accelerator.synchronize()
    • Accelerator.tpu
    • Accelerator.trainium
  • KernelBackend
    • KernelBackend.cuda
    • KernelBackend.get_compatible_accelerator()
    • KernelBackend.mps
    • KernelBackend.nki
    • KernelBackend.pallas
    • KernelBackend.rocm
    • KernelBackend.torch
    • KernelBackend.triton
    • KernelBackend.verify_accelerator()
• xma.autotuner
  • xma.autotuner.config
    • AutotuneConfig
      • AutotuneConfig.get_key_values()
      • AutotuneConfig.is_condition_valid()
    • get_cartesian_product_autotune_configs()
  • xma.autotuner.tuner
    • AutotunedFunction
      • AutotunedFunction.exposed_signature
    • autotune()
• xma.counters
  • enable_counters()
  • get_counter_value()
  • increment_counter()
  • is_counter_enabled()
  • reset_counters()
• xma.custom_op
  • CustomOp
    • CustomOp.backward()
    • CustomOp.forward()
    • CustomOp.forward_backward_torch()
    • CustomOp.run()
  • ctx_needs_gradients()
  • ctx_save_for_backward()
  • xma_op()
• xma.functional
  • xma.functional.bmm
    • bmm()
  • xma.functional.continuous_count
    • continuous_count()
  • xma.functional.cross_entropy
    • cross_entropy()
  • xma.functional.fused_linear_cross_entropy
    • fused_linear_cross_entropy()
  • xma.functional.fused_residual_add_rmsnorm
    • fused_residual_add_rmsnorm()
    • fused_residual_add_rmsnorm_backward()
    • fused_residual_add_rmsnorm_forward()
  • xma.functional.p_norm
    • p_norm()
  • xma.functional.rmsnorm
    • rmsnorm()
  • xma.functional.sequence_packing
    • pack_sequence()
    • unpack_sequence()
  • xma.functional.softmax
    • softmax()
  • xma.functional.swiglu
    • swiglu()
    • swiglu_packed()
• xma.inductor
  • enable_kernels()
  • get_fused_residual_add_rmsnorm_replacer()
  • get_rmsnorm_replacer()
  • partialize_and_update_signature()
• xma.jit
  • cpp_jit()
• xma.layers
  • xma.layers.gru
    • GRU
      • GRU.extra_repr()
      • GRU.forward()
      • GRU.reset_parameters()
    • gru()
  • xma.layers.linear_attention
    • LinearAttention
      • LinearAttention.forward()
    • linear_attention()
  • xma.layers.m2rnn
    • M2RNN
      • M2RNN.forward()
      • M2RNN.reset_parameters()
    • m2rnn()
  • xma.layers.moe
    • Experts
      • Experts.down_projection_triton_forward()
      • Experts.extra_repr()
      • Experts.reset_parameters()
      • Experts.torch_forward()
      • Experts.up_projection_triton_forward()
    • MoE
      • MoE.forward()
  • xma.layers.rnn
    • RNN
      • RNN.extra_repr()
      • RNN.forward()
      • RNN.reset_parameters()
    • rnn()
• xma.math
  • ceil_divide()
  • check_power_of_2()
  • divide_if_divisible()
  • get_next_power_of_2()
  • get_powers_of_2()
• xma.optimizers
  • xma.optimizers.sgd
    • SGD
      • SGD.step()
    • sgd()