vahidzee / dypy

Toolkit for Dynamic Python code manipulations

• 👋 Hi, I’m Vahid Zehtab (@vahidzee)
• 👀 I’m interested in investigative studies in machine learning, especially for generative models and probabilistic deep learning models.
• 📫 Feel free to reach me at [email protected]

dypy's Issues

Adding context to class wrappers

Sometimes we need to add a context to a class wrapper for calling it later on from other sections of the code:

# context can be a list of multiple contexts to consider
@dyw.dynamize(add_contexts=['torch'])
class A:
    layer = dyw.composite('nn.Linear`, in_features=10, out_features=100)

pass down in composition

Check out the following code that we want to dynamize:

import typing as th
import torch
from .linear import MaskedLinear
import dycode


class MaskedBlock(MaskedLinear):
    def __init__(
        self,
        # linear args
        in_features: th.Union[th.List[int], int],
        out_features: th.Union[th.List[int], int],
        bias: bool = True,
        # activation
        activation: th.Optional[str] = None,
        activation_args: th.Optional[dict] = None,
        # batch norm
        batch_norm: bool = False,
        batch_norm_args: th.Optional[dict] = None,
        # ordering args
        auto_connection: bool = True,
        # general parameters
        device: th.Optional[torch.device] = None,
        dtype: th.Optional[torch.dtype] = None,
        mask_dtype: torch.dtype = torch.uint8,
    ):
        # init linear
        super().__init__(
            in_features=in_features,
            out_features=out_features,
            bias=bias,
            auto_connection=auto_connection,
            device=device,
            dtype=dtype,
            mask_dtype=mask_dtype,
        )

        self.activation = dycode.eval(activation)(**(activation_args or dict())) if activation else None
        self.batch_norm = (
            torch.nn.BatchNorm1d(num_features=out_features, dtype=dtype, device=device, **(batch_norm_args or dict()))
            if batch_norm
            else None
        )

    def forward(self, inputs: torch.Tensor, perm_mat: torch.Tensor) -> torch.Tensor:
        outputs = super().forward(inputs, perm_mat)
        outputs = self.activation(outputs) if self.activation else outputs
        outputs = self.batch_norm(outputs) if self.batch_norm else outputs
        return outputs

Note that if we rewrite batch_norm as a dyw.composite then the device and dtype of the composite objects might defer from the whole object.

Therefore, it is good to have some pass_down arguments; for example, MaskedBlock(pass_down_device=sth) not only will it set device to sth for masked block, but also set all the composite object device attributes to sth as well.

Syntax upgrade for Inheritance and Composite wrappers

In addition to suggestions on improving the proposed syntax for dynamizing Inheritance, I suggest we generalize everything into a unified methodology of parsing the arguments passed on to the __init__ of a class that has been decorated with dycode.dynamize.

As a showcase, imagine the code bellow:

@dy.dynamize(
    dynamize_init=True, 
    extends=[torch.nn.Linear], 
    merge_args=dict(
        in_features=["conv__in_channels", "torch.nn.Linear__in_features"],
        device=None,
        dtype=None,
     ) # or perhaps extends_merge_args=["dytpe", "device"] to only share those
 )
class LinearBlock(torch.nn.Linear):
   activation: th.Optional[torch.nn.Module] = dy.composite(default="torch.nn.ReLU", required=False, null_value=None)
   batch_norm: th.Optional[torch.nn.Module] = dy.composite(default="torch.nn.BatchNorm1d, required=False)
   conv= dy.composite(torch.nn.Conv1d)
   
   def __init__(self):
        # a linear block, with optional activation and batchnorm, which passes the results to a conv1d as well
        super().__init__(...) # instantiate the linear layer
        ....

Although the example might not be sound in any conventional DL model, similar scenarios regularly arise when developing deep models.

A neat feature that dycode could potentially handle would be to infer what arguments are not present in the implementation of original LinearBlock's __init__, and how to pass composite arguments or to share parameters between them and the inherited classes.

For instance, here, except for usual implementations of activation, all batch_norm, conv, and even the base Linear modules accept a device or a dtype argument at initialization. Mentioning "dtype" and "device" in the extends_merge_args, should tell dycode to add a single device and dtype argument to LinearBlock, so that we can manipulate those at initialization of LinearBlock by LinearBlock(dtype=sth, device=sth).

Or that, if a parameter could be shared among composite classes and/or base classes but has different names in their constructors, we can mention a shared argument name and what those arguments translate to in merge_args.
Of course, many things here can be automatically inferred, and it might also make sense for dycode.dynamize to automatically parse the init in this way by default.

Another helpful feature would be the ability to manipulate arguments passed on to composite/base classes from the dynamized constructor using the same __ syntax. For instance, instead of providing args for each composite class, why not just parse the arguments passed on to the constructor and see where each provided parameter should go:

LinearBlock(in_features=2, activation__in_place=True, device='cpu', kernel_size=3, conv__padding=0)

In this example, it's obvious where each parameter should go, and it's pretty readable too.
Why not automate the same idea?
Stack all the signatures of composite and inherited base classes, try to merge their arguments into one single translation table, use the merge_args to merge those which could be shared, or maybe merge them by default for those with the same name if merge_args=True, then use such a table to dynamically evaluate the arguments passed on to the dynamized class.

Add unittests

All the tests are now in the ipynb file. One should create a test for unittesting.

Enhance dypy.method wrapper

Add essentials like the ability to combine with @property or introduce shortcuts for managing properties.

Make dynamically processed args accessible to __init__.

The current implementation of dynamized __init__ separates the context for calling the original __init__ implementation, it would be very helpful to mix the two contexts so that arguments which will be added to the init later on by dycode are also accessible.

For instance:

@dy.dynamize
class MyClass:
   x = dy.field()
   def __init__(self, a, b, c, ...):
      # doing something with x should also be possible

Adding Nested Composite Instantiation

Many times, the following scenario might happen:

class A:
    def __init__(self, x, y) -> None:
        self.x = x
        self.y = y

class B:
    def __init__(self, a, b) -> None:
        self.a = a
        self.b = b

class C:
    def __init__(self, c, d, a, b, x, y) -> None:
        self.c = c
        self.d = d
        self.sample_a = A(a, b)
        self.sample_b = B(x, y)

where sample_a and sample_b are composite objects of C, when the composition becomes deeper and deeper, all of the values in all the attributes of these composite variables should be passed on to the constructor, making it far less readable. We can fix this using the following functionality, which defines a set of composite fields.

Then while instantiating the end class, we can either leave the composite objects as is using the default setting, or we can customize their values in the instructor. Check out the following piece of code which contains nested instantiations:

import dycode.wrappers as dyw

class A:
    def __init__(self, a, b):
        self.a = a
        self.b = b

@dyw.dynamize
class B_P:
    sample_a = dyw.composite('A', a=-1, b=-1)
    def __init__(self, xp):
        self.xp = xp

@dyw.dynamize
class B:
    sample_bp = dyw.composite('B_P', xp=-1)
    def __init__(self, x):
        self.x = x

@dyw.dynamize
class C:
    sample_a = dyw.composite('A', a = 2, b = 5)
    def __init__(self, a):
        self.a = a

# a simple constructor
sample1 = C(10)
print(sample1.sample_a.a)
print("--------------------")
# A little configuration
sample1_5 = C(10, sample_a_args={'b': 101})
print(sample1_5.sample_a.a)
print(sample1_5.sample_a.b)
print("--------------------")
# changing a deeper field in the constructor
sample2 = C(10, sample_a='B', sample_a_args={'x': 20})
print(sample2.sample_a.x)
print("--------------------")
# Changing even deeper field in the constructor
sample3 = C(11, 
    sample_a='B',
    sample_a_args={
        'x': 22, 
        'sample_bp': 'B_P', 
        'sample_bp_args': {
            'xp': 33
        }
    }
)
print(sample3.a)
print(sample3.sample_a.x)
print(sample3.sample_a.sample_bp.xp)
print("--------------------")

This can particularly help in lazy configurations. For example, detectron2 has implemented this here, however, when decorating any class you can seamlessly do this on top of any dataclass.

Fast Inheritted Instantiation

Just like Composites, sometimes when something is inherited then all the way down in the inheritance tree everything should be included in the constructors.

class A:
    def __init__(self, x, y) -> None:
        self.x = x
        self.y = y
class B:
    def __init__(self, a, b) -> None:
        self.a = a
        self.b = b

class C(A, B):
    def __init__(self, c, d, a, b, x, y) -> None:
        self.c = c
        self.d = d
        A.__init__(self, x, y)
        B.__init__(self, a, b)

class D(C):
    def __init__(self, e, f, c, d, a, b, x, y) -> None:
        super().__init__(c, d, a, b, x, y)
        self.e = e
        self.f = f

class E(D):
    def __init__(self, g, h, e, f, c, d, a, b, x, y) -> None:
        super().__init__(e, f, c, d, a, b, x, y)
        self.g = g
        self.h = h

In that case, we can use a notion of dyw.inherited_field as follows:

class A:
    def __init__(self, x, y) -> None:
        self.x = x
        self.y = y
class B:
    def __init__(self, a, b) -> None:
        self.a = a
        self.b = b

@dyw.dynamize
class C(A, B):
    x = dyw.inherited_field(default=1, parent=A)
    y = dyw.inherited_field(default=2, parent=A)
    
    a = dyw.inherited_field(default=10, parent=B)
    b = dyw.inherited_field(default=20, parent=B)

    c = dyw.field(default=100)
    d = dyw.field(default=200)

@dyw.dynamize
class D(C):
    e = dyw.field(default=11)
    f = dyw.field(default=22)

@dyw.dynamize
class E(D):
    g = dyw.field(default=111)
    h = dyw.field(default=222)

This should then be similarly instantiated as the previous one.

Make fields Nullable

For any dynamic field, we can add an option to make it nullable. For example,

@dyw.dynamize
class TrainingModule:
    layer = dyw.composite('torch.nn.Linear`, nullable=True)
    a: int = dyw.field(1, nullable=True)

Now let's instantiate:

tm1 = TrainingModule(layer=None, a=2)
tm2 = TrainingModule(layer='torch.nn.MaxPool1d', a=None)