Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers