Updated optimize to handle combination of multiple intervention templates

[anirban-chaudhuri]

• created a lambda function to combine different intervention templates
• updated type hinting for static_parameter_interventions input to optimize to accommodate lambda functions
• downstream changes in ouu.py to handle lambda functions

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[anirban-chaudhuri]

@SamWitty Any ideas on how to make this cleaner and if there is a way to create a generic template for combining the different intervention builders into a lambda function?

• Currently the type hint has become a bit complicated for interventions in optimize and we might want to think of a way to simplify it, if possible
• The end of the optimize_interface has the lambda function implementation and the downstream changes required to accommodate that change in ouu.py

[SamWitty]

I think this PR addresses the issue, but it would be nice to have the explicit intervention template combinator provided (which I believe is now commented out) instead of combining them manually in a lambda function. This might be worth hopping on a chat and working through together.

[SamWitty]

This is fine. If you want you can make this a bit more concise with a type alias.

E.g. https://github.com/BasisResearch/chirho/blob/master/chirho/dynamical/ops.py#L14

Here, you could add a line with the following:

InterventionFunc = Callable[[torch.Tensor], Dict[float, Dict[str, Intervention]]]

Then this type could be:

Union[InterventionFunc, Callable[[torch.Tensor], List[InterventionFunc]]

[SamWitty]

Looking again, is this type signature correct? Should it be:

Callable[[torch.Tensor], Dict[float, Dict[str, Intervention]]],
        Callable[
            [torch.Tensor],
            List[Dict[float, Dict[str, Intervention]]],
        ]

[anirban-chaudhuri]

I am not sure. It might be. I was thinking once I create the lambda function it becomes a Callable of Callable functions.. Let's discuss when we meet.

[SamWitty]

I do think it would be good to provide this instead of the manual lambda function expansion you have in the tests.

[anirban-chaudhuri]

I completely agree. I couldn't write a bug free one so I left it. We can hop on a call to make it happen.

[SamWitty]

Looks good. One minor nitpick.

I think having the interventions be mapping from float to interventions is the right idea. We can make the interfaces.py consistent with this in a separate PR if they've drifted.

[SamWitty]

Could you remove this commented out code?

[anirban-chaudhuri]

Done

[anirban-chaudhuri]

Looks good. One minor nitpick.

I think having the interventions be mapping from float to interventions is the right idea. We can make the interfaces.py consistent with this in a separate PR if they've drifted.

Sounds good. I will leave my comment from slack here so that we have the specific details.

Type hinting needs to be fixed for interventions: Should it be Dict[float, Dict[str, Intervention]] or Dict[torch.Tensor, Dict[str, Intervention]] ? Both are used in the interfaces.py.