[flax] Widen LogicalNames type to accept multi-axis tuples

* Widen `flax.typing.LogicalNames` to `tuple[AxisName, ...]`.
* Replace redundant local `LogicalNames` definitions in videogvt with `nn.meta.LogicalNames`.

PiperOrigin-RevId: 973968118

Author: adityagupta1089

flax/typing.py

@@ -124,8 +124,8 @@ class Out(Generic[T]):
 
 # SPMD
 
-LogicalNames = tuple[Union[str, None], ...]
 AxisName = str | tuple[str, ...] | None
+LogicalNames = tuple[AxisName, ...]
 
 # Maps each logical axis  to physical mesh, can be either None (replicated),
 # one physical axis or a tuple of physical axes.