Stage "handles" missing in HCL models created via classes.

[jcasas00]

For code management/organization reasons, we started using classes to create HCL models. For simply creating the schedule, building and running the model, it seems to be working fine.
However, I noticed now that references to "Stage" objects are missing when models are defined this way.

For example (using the two_stage example in the documentation), when defined the "normal" way:

def two_stage(A):
    B = hcl.compute(A.shape, lambda x, y: A[x, y] + 1, "B")
    C = hcl.compute(A.shape, lambda x, y: B[x, y] + 1, "C")
    return C

I can refer to stage "B" as two_stage.B. And indeed I see attributes 'B' and 'C':

>>> dir(two_stage)
['B', 'C', '__annotations__', '__call__', '__class__', '__closure__', '__code__', '__defaults__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__get__', '__getattribute__', '__globals__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__kwdefaults__', '__le__', '__lt__', '__module__', '__name__', '__ne__', '__new__', '__qualname__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__']
>>> type (two_stage.B)
<class 'heterocl.schedule.Stage'>

But if we define it in a class like:

class testme:
    def __init__(self):
        pass

    def two_stage(self, A):
        B = hcl.compute(A.shape, lambda x, y: A[x, y] + 1, "B")
        C = hcl.compute(A.shape, lambda x, y: B[x, y] + 1, "C")
        return C

    @staticmethod
    def two_stage_static(A):
        B = hcl.compute(A.shape, lambda x, y: A[x, y] + 1, "B")
        C = hcl.compute(A.shape, lambda x, y: B[x, y] + 1, "C")
        return C

I no longer see the 'B' and 'C' attributes. This is true for both two_stage (normal method) and two_stage_static (which should look like a regular function).

>>>dir(testme.two_stage)
['__annotations__', '__call__', '__class__', '__closure__', '__code__', '__defaults__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__get__', '__getattribute__', '__globals__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__kwdefaults__', '__le__', '__lt__', '__module__', '__name__', '__ne__', '__new__', '__qualname__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__']
>>>dir(testme.two_stage_static)
['__annotations__', '__call__', '__class__', '__closure__', '__code__', '__defaults__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__get__', '__getattribute__', '__globals__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__kwdefaults__', '__le__', '__lt__', '__module__', '__name__', '__ne__', '__new__', '__qualname__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__']

Any recommendations/ideas ?

[pasqoc]

The following seems to work (note the wrapper for the member function version):

A = hcl.placeholder((10,10), "A")


def two_stage(A):
    B = hcl.compute(A.shape, lambda x, y: A[x, y] + 1, "B")
    C = hcl.compute(A.shape, lambda x, y: B[x, y] + 1, "C")
    return C

class testme:
    def __init__(self):
        pass

    def two_stage(self, A):
        B = hcl.compute(A.shape, lambda x, y: A[x, y] + 1, "B")
        C = hcl.compute(A.shape, lambda x, y: B[x, y] + 1, "C")
        return C

    @staticmethod
    def two_stage_static(A):
        B = hcl.compute(A.shape, lambda x, y: A[x, y] + 1, "B")
        C = hcl.compute(A.shape, lambda x, y: B[x, y] + 1, "C")
        return C


t = testme()
def two_stage_member(*i):
  t.two_stage(*i)

s1 = hcl.create_schedule([A], two_stage)
dir(two_stage)

s2 = hcl.create_schedule([A], testme.two_stage_static)
dir(testme.two_stage_static)

s3 = hcl.create_schedule([A], two_stage_member)
dir(t.two_stage)

[seanlatias]

The members are associated when hcl.create_schedule is called. And the target to be associated is the function that we use as the argument, as shown in @pasqoc's solution.