[sharding_in_types] Enforce AxisTypes to always exist if `set_mesh` i… #68019

Job	Run time
FFI example	2m 29s
Documentation - test code snippets (3.10)	2m 6s
build with 3.10 (py 3.10 on ubuntu-20.04, x64=1)	10m 0s
Documentation - render documentation (3.10)	3m 18s
jax2tf_test (py 3.10 on ubuntu-latest, x64=0)	1m 24s
lint_and_typecheck	1m 4s
build with 3.13 (py 3.13 on ubuntu-20.04, x64=0)	5m 22s
	25m 43s

Provide feedback