from lapjax._src.lax.control_flow import conditionals as conditionals
from lapjax._src.lax.control_flow import solves as solves
from lapjax._src.lax.control_flow import for_loop as for_loop
from lapjax._src.lax.control_flow import common as common
from lapjax._src.lax.control_flow import loops as loops
import sys, importlib
from lapjax.lapsrc.wrapper import _wrap_module
_wrap_module(importlib.import_module(__name__.replace('lapjax', 'jax')), 
             sys.modules[__name__])
