Format docstrings in utils/insert_along_axis.py