brainstate 0.1.0.post20241219__py2.py3-none-any.whl → 0.1.0.post20241220__py2.py3-none-any.whl

brainstate/compile/_jit.py

@@ -40,7 +40,7 @@ class JittedFunction(Callable):
     jitted_fun: jax.stages.Wrapped  # the jitted function
     clear_cache: Callable  # clear the cache of the jitted function
     eval_shape: Callable  # evaluate the shape of the jitted function
-    lower: Callable  # lower the jitted function
+    compile: Callable  # lower the jitted function
     trace: Callable  # trace the jitted
 
     def __call__(self, *args, **kwargs):
@@ -104,7 +104,18 @@ def _get_jitted_fun(
     def eval_shape():
         raise NotImplementedError
 
-    def lower():
+    def trace():
+        """Trace this function explicitly for the given arguments.
+
+        A traced function is staged out of Python and translated to a jaxpr. It is
+        ready for lowering but not yet lowered.
+
+        Returns:
+          A ``Traced`` instance representing the tracing.
+        """
+        raise NotImplementedError
+
+    def compile(*args, **params):
         """Lower this function explicitly for the given arguments.
 
         A lowered function is staged out of Python and translated to a
@@ -114,18 +125,13 @@ def _get_jitted_fun(
         Returns:
           A ``Lowered`` instance representing the lowering.
         """
-        raise NotImplementedError
-
-    def trace():
-        """Trace this function explicitly for the given arguments.
+        # compile the function and get the state trace
+        state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
+        read_state_vals = state_trace.get_read_state_values(True)
 
-        A traced function is staged out of Python and translated to a jaxpr. It is
-        ready for lowering but not yet lowered.
+        # call the jitted function
+        return jit_fun.lower(state_trace.get_state_values(), *args, **params).compile()
 
-        Returns:
-          A ``Traced`` instance representing the tracing.
-        """
-        raise NotImplementedError
 
     jitted_fun: JittedFunction
 
@@ -144,8 +150,8 @@ def _get_jitted_fun(
     # evaluate the shape of the jitted function
     jitted_fun.eval_shape = eval_shape
 
-    # lower the jitted function
-    jitted_fun.lower = lower
+    # compile the jitted function
+    jitted_fun.compile = compile
 
     # trace the jitted
     jitted_fun.trace = trace

{brainstate-0.1.0.post20241219.dist-info → brainstate-0.1.0.post20241220.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20241219
+Version: 0.1.0.post20241220
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

{brainstate-0.1.0.post20241219.dist-info → brainstate-0.1.0.post20241220.dist-info}/RECORD

@@ -24,7 +24,7 @@ brainstate/compile/_conditions.py,sha256=gApsHKGQrf1QBjoKXDVL7VsoeJ2zFtSc-hFz9nb
 brainstate/compile/_conditions_test.py,sha256=s9LF6h9LvigvgxUIugTqvgCHBIU8TXS1Ar1OlIxXfrw,8389
 brainstate/compile/_error_if.py,sha256=TFvhqITKkRO9m30GdlUP4eEjJvLWQUhjkujXO9zvrWs,2689
 brainstate/compile/_error_if_test.py,sha256=SJmAfosVoGd4vhfFtb1IvjeFVW914bfTccCg6DoLWYk,1992
-brainstate/compile/_jit.py,sha256=3mQ-RUFz35wceZyKE_MoR58OBL0RK_i6sHm4rWYzMLs,13698
+brainstate/compile/_jit.py,sha256=3WBXNTALWPYC9rQH0TPH6w4bjG0BpnZt3RAzUQF5kkc,14045
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
 brainstate/compile/_loop_collect_return.py,sha256=_iOVPytbctgyaIOxQZH3A2ZbsSoT7VXnFk6Q6R8-gvA,23360
 brainstate/compile/_loop_collect_return_test.py,sha256=bA-_11E8A_0jR5umEO3e409y7bb5QYDTgSL-SBaX7kQ,1802
@@ -137,8 +137,8 @@ brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7
 brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
 brainstate/util/_tracers.py,sha256=0r5T4nhxMzI79NtqroqitsdMT4YfpgV5RdYJLS5uJ0w,2285
 brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.post20241219.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20241219.dist-info/METADATA,sha256=nOo68Iuy-6B4po__cxVsBZlUEKbj0tZ1n0gq6x7MqLk,3533
-brainstate-0.1.0.post20241219.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20241219.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20241219.dist-info/RECORD,,
+brainstate-0.1.0.post20241220.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20241220.dist-info/METADATA,sha256=lfkUbD1vYx4bikkUolrL_pCp4zryfvGExXyWP_QH5tM,3533
+brainstate-0.1.0.post20241220.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20241220.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20241220.dist-info/RECORD,,