https://github.com/patrick-kidger/sympy2jax/blob/main/sympy2jax/sympy_module.py taken from https://github.com/scikit-hep/pyhf/pull/2579#discussion_r2001115783
