brainpy / BrainPy

Brain Dynamics Programming in Python
https://brainpy.readthedocs.io/
GNU General Public License v3.0
531 stars 94 forks source link

Support to trace ``Variable``s inside of a brainpy transformation #471

Closed chaoming0625 closed 1 year ago

chaoming0625 commented 1 year ago

Please:

The current brainpy transformations can only trace Variable outside these transformation functions. For example,


class A():
  def __init__(self):
     self.a  = bm.Variable(1)  # can be traced

  @bm.cls_jit
  def f(self):
     self.b = bm.Variable(1)   # cannot be traced
     self.a += self.b

This is in contrast to the OO transformations in Flax, which supports to trace Variables inside of a transformation function.

It may be essential to support Variable tracing during a computation.

chaoming0625 commented 1 year ago

Has been solved in #472.