Fix megablox tile sizes #11155

Job	Run time
Analyze Code Changes for Test Orchestration	19s
build_and_upload_maxtext_package / build_and_upload	38s
maxtext_tpu_pathways_unit_tests / run	18m 27s
All Notebooks Passed	3s
maxtext_tpu_pathways_integration_tests / run	18m 27s
gpu-integration tests / gpu-integration / run	2m 36s
cpu-post-training-unit tests / cpu-post-training-unit / run	5m 22s
gpu-unit tests / gpu-unit / run	18m 27s
cpu-unit tests / cpu-unit / run	17m 29s
cpu-unit tests / cpu-unit / run	4m 18s
tpu-integration tests / tpu-integration / run	12m 10s
tpu-unit tests / tpu-unit / run	18m 27s
maxtext_jupyter_notebooks	-1s
tpu-post-training-unit tests / tpu-post-training-unit / run	18m 27s
Notify failed build	2s
All Required Tests Passed	3s
	2h 15m 14s

Provide feedback