Replace uses of jax.partial with functools.partial, in preparation for removing jax.partial.
Replace uses of jax.partial with functools.partial, in preparation for removing jax.partial.

jax.partial is an alias for functools.partial, and functools.partial is a Python standard library API. There's no need for jax to export this function.

PiperOrigin-RevId: 395993004

HopkinsIDD/EpiForecastStatMech,HopkinsIDD/EpiForecastStatMech