add short hand notation for node addition without dataclass field declaration about pytreeclass HOT 1 CLOSED

@treeclass
class StackedLinear:

    def __init__(self, key):
        self.keys = jax.random.split(key, 3)

    def __call__(self, x):
        x = self.register_node(Linear(key=self.keys[0], in_dim=x.shape[-1], out_dim=128),name="l1")(x)
        x = jax.nn.tanh(x)
        x = self.register_node(Linear(key=self.keys[1], in_dim=128, out_dim=128),name="l2")(x)
        x = jax.nn.tanh(x)
        x = self.register_node(Linear(key=self.keys[2], in_dim=128, out_dim=x.shape[-1]),name="l3")(x)

        return x

@treeclass
class StackedLinear:
    l1: Linear
    l2: Linear
    l3: Linear

    def __init__(self, key, in_dim, out_dim):

        keys = jax.random.split(key, 3)

        self.l1 = Linear(key=keys[0], in_dim=in_dim, out_dim=128)
        self.l2 = Linear(key=keys[1], in_dim=128, out_dim=128)
        self.l3 = Linear(key=keys[2], in_dim=128, out_dim=out_dim)

    def __call__(self, x):
        x = self.l1(x)
        x = jax.nn.tanh(x)
        x = self.l2(x)
        x = jax.nn.tanh(x)
        x = self.l3(x)

        return x