fix(jax): support charge-spin savedmodel export#5737

njzjz wants to merge 3 commits into

deepmodeling:masterdeepmodeling/deepmd-kit:masterfrom

njzjz:fix/jax-charge-spin-savedmodelnjzjz/deepmd-kit:fix/jax-charge-spin-savedmodelCopy head branch name to clipboard