Currently, we expose the APIs in the way that prioritizes intervention serializations, which means we can only use intervention primitives defined in pyvene which limits customizable intervention types. Although people can define interventions as nn.Modules, it takes more efforts comparing to a simple function, or lambda expression.
This change introduces how to interact with pyvene with a one-line lambda expression-based intervention which can also do trainable interventions with more customizable formats. At this point, pyvene basically acts like a torch.hook manager only losing its ability to share and save. Yet, it can easily support those simple interventions that are designed to be customizable.
Testing Done
grad_fn=<AddBackward0>)
./juice/scr/wuzhengx/pyvene/pyvene/models/interventions.py:421: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
mask_sigmoid = torch.sigmoid(self.mask / torch.tensor(self.temperature))
.........Removing testing dir ./test_output_dir_prefix-2891ac
Removing testing dir ./test_output_dir_prefix-169855
Removing testing dir ./test_output_dir_prefix-6f0ab4
Removing testing dir ./test_output_dir_prefix-be15ba
Removing testing dir ./test_output_dir_prefix-50bb90
.............
----------------------------------------------------------------------
Ran 63 tests in 44.468s
OK
Added in tutorials in pyvene-101.
Checklist:
[x] My PR title strictly follows the format: [Your Priority] Your Title
Description
Currently, we expose the APIs in the way that prioritizes intervention serializations, which means we can only use intervention primitives defined in pyvene which limits customizable intervention types. Although people can define interventions as
nn.Modules
, it takes more efforts comparing to a simple function, or lambda expression.This change introduces how to interact with pyvene with a one-line lambda expression-based intervention which can also do trainable interventions with more customizable formats. At this point, pyvene basically acts like a
torch.hook
manager only losing its ability to share and save. Yet, it can easily support those simple interventions that are designed to be customizable.Testing Done
Checklist:
[Your Priority] Your Title