feat(jax): support neural networks#4156

wanghan-iapcm merged 3 commits into

deepmodeling:develdeepmodeling/deepmd-kit:develfrom

njzjz:jax-startnjzjz/deepmd-kit:jax-startCopy head branch name to clipboard