xma

• xma.accelerator
  • Accelerator
    • Accelerator.cpu
    • Accelerator.cuda
    • Accelerator.device_count()
    • Accelerator.get_accelerator()
    • Accelerator.get_core_count()
    • Accelerator.get_current_device()
    • Accelerator.get_kernel_backend()
    • Accelerator.get_lane_count()
    • Accelerator.mps
    • Accelerator.rocm
    • Accelerator.synchronize()
    • Accelerator.tpu
    • Accelerator.trainium
  • KernelBackend
    • KernelBackend.cuda
    • KernelBackend.get_compatible_accelerator()
    • KernelBackend.jax
    • 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
• xma.counters
  • enable_counters()
  • get_counter_value()
  • increment_counter()
  • is_counter_enabled()
  • reset_counters()
• xma.custom_op
• 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.inductor
• xma.jit
• 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.math
  • ceil_divide()
  • check_power_of_2()
  • divide_if_divisible()
  • get_next_power_of_2()
  • get_powers_of_2()
• xma.optimizers
  • xma.optimizers.sgd
    • xma.optimizers.sgd.param_group