fix(jax): workaround for "xxTracer is not a valid JAX type"#4776

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

njzjz:jax_workaroundnjzjz/deepmd-kit:jax_workaroundCopy head branch name to clipboard