if "a" <= letter <= "z" or letter == "_":
elementary_axes.append(letter)
else:
assert "0" <= letter <= "9"
prev_letter = composite_axis_name.pop()
assert "a" <= prev_letter <= "z"
elementary_axes.append(prev_letter + letter)
return elementary_axes[::-1]
After Change
def get_axes_names(composite_axis_name: str):
axes_names = [axis for axis in composite_axis_name.split(" ") if len(axis) > 0]
for axis in axes_names:
if axis == "_":
continue
assert "a" <= axis[0] <= "z"
for letter in axis:
assert str.isdigit(letter) or "a" <= letter <= "z"
return axes_names
