feat(jax): add training#5460

njzjz merged 6 commits into

deepmodeling:masterdeepmodeling/deepmd-kit:masterfrom

njzjz-bothub:cleanup/parallel-no-hessiannjzjz-bothub/deepmd-kit:cleanup/parallel-no-hessianCopy head branch name to clipboard