maxtext.kernels.megablox package

Megablox kernel

Submodules

• maxtext.kernels.megablox.backend module
  • make_group_metadata()
  • gmm()
  • tgmm()
• maxtext.kernels.megablox.common module
  • is_tpu()
  • tpu_kind()
  • tpu_generation()
  • supports_bfloat16_matmul()
  • assert_is_supported_dtype()
  • select_input_dtype()
• maxtext.kernels.megablox.ops module
  • gmm()