chore: use `xp.take_along_axis` if Array API version >=2024.12#4406

wanghan-iapcm merged 1 commit intodeepmodeling:develdeepmodeling/deepmd-kit:develfrom

njzjz:xp-take-along-axisnjzjz/deepmd-kit:xp-take-along-axisCopy head branch name to clipboard