remove jax from import

[OzanCatalVerses]

si.py was still refering to the pymdp.jax.control version of control.py breaking imports of sophisticated inference. Also removes unused imports.

[tverbele]

Yup, good catch ...
@conorheins do we have a new equivalent of the old pymdp.maths.EPS_VAL instead of hard coding 1e-16?
We might also want to update the notebooks though...

[conorheins]

Yup, good catch ... @conorheins do we have a new equivalent of the old pymdp.maths.EPS_VAL instead of hard coding 1e-16? We might also want to update the notebooks though...

See my comment above on the policy entropy function -- we avoid evaluating log of 0.0 within the the stablelog function now, because that function clips the output to a tiny value determined by jnp.finfo(float).eps

[conorheins]

Can you replace that call to jnp.log(node["q_pi"] + 1e-16) with the stablelog call?

[conorheins]

This may be fixable by simply doing

-jnp.dot(node["q_pi"], pymdp.maths.stablelog(node["q_pi"]))`

or even

-pymdp.maths.stable_xlogx(node["q_pi"]).sum()