fix(jax): calculate virial in `call_lower`#4304

wanghan-iapcm merged 2 commits intodeepmodeling:develdeepmodeling/deepmd-kit:develfrom

njzjz:jax-call-lower-virialnjzjz/deepmd-kit:jax-call-lower-virialCopy head branch name to clipboard