from fjformer.pallas_operations.tpu.flash_attention.flash_attention import (
	BlockSizes as BlockSizes,
)
from fjformer.pallas_operations.tpu.flash_attention.flash_attention import (
	flash_attention as flash_attention,
)
