Add jax.numpy.linalg (#316)

* add jax.numpy.linalg

* fix relative import

---------

Co-authored-by: Eli <eli@elibingham.com>

Author: jfeser

effectful/handlers/jax/numpy.py effectful/handlers/jax/numpy/__init__.pyeffectful/handlers/jax/numpy.py renamed to effectful/handlers/jax/numpy/__init__.py

@@ -2,7 +2,7 @@
 
 import jax.numpy
 
-from ._handlers import _register_jax_op, _register_jax_op_no_partial_eval
+from .._handlers import _register_jax_op, _register_jax_op_no_partial_eval
 
 _no_overload = ["array", "asarray"]
 

effectful/handlers/jax/numpy/linalg.py

@@ -0,0 +1,13 @@
+from typing import TYPE_CHECKING
+
+import jax.numpy.linalg
+
+from effectful.handlers.jax._handlers import _register_jax_op
+
+for name, op in jax.numpy.linalg.__dict__.items():
+    if callable(op):
+        globals()[name] = _register_jax_op(op)
+
+# Tell mypy about our wrapped functions.
+if TYPE_CHECKING:
+    from jax.numpy.linalg import *  # noqa: F403