ejkernel.kernels._pallas.tpu.grouped_matmul package

• ejkernel.kernels._pallas.tpu.grouped_matmul._interface
  • grouped_matmul()
• ejkernel.kernels._pallas.tpu.grouped_matmul._pallas_impl
  • grouped_matmul()
  • make_group_metadata()
  • transposed_grouped_matmul()
• ejkernel.kernels._pallas.tpu.grouped_matmul._utils
  • assert_is_supported_dtype()
  • is_tpu()
  • select_input_dtype()
  • supports_bfloat16_matmul()
  • tpu_generation()
  • tpu_kind()