Add module for jax.scipy.special (#294)

* add module for jax.scipy.special

* lint

---------

Co-authored-by: eb8680 <eb8680@users.noreply.github.com>

Author: jfeser

effectful/handlers/jax/scipy/special.py

@@ -0,0 +1,11 @@
+from typing import TYPE_CHECKING
+
+import jax.scipy.special
+
+from effectful.handlers.jax._handlers import _register_jax_op
+
+logsumexp = _register_jax_op(jax.scipy.special.logsumexp)
+
+# Tell mypy about our wrapped functions.
+if TYPE_CHECKING:
+    from jax.scipy.special import logsumexp  # noqa: F401