maxtext.kernels.megablox.common module#
Common utilities for GMM kernels.
- maxtext.kernels.megablox.common.tpu_kind()[source]#
Query identification string for the currently attached TPU.
- Return type:
str
- maxtext.kernels.megablox.common.tpu_generation()[source]#
Generation number of the currently attached TPU.
- Return type:
int
- maxtext.kernels.megablox.common.supports_bfloat16_matmul()[source]#
Does the currently attached CPU support bfloat16 inputs?
- Return type:
bool