aai-institute / aai-institute/pyDVL
Use `pytree` operations to simplify code in influence
- Dominant language
- Python
- Stars
- 146
- Forks
- 10
- PR merge metrics
- No merged PRs in 30d
Description
## What are pytrees?
While reviewing #582, I got the feeling that we have quite a bit of code that re-implements the typical `pytree` operations as implemented in [jax](https://jax.readthedocs.io/en/latest/jax.tree_util.html), [optree](https://github.com/metaopt/optree?tab=readme-ov-file#tree-map-nargs) or my instructional library [pybaum](https://github.com/OpenSourceEconomics/pybaum).
The pytree operations solve the following problem in a general way: In math notation, we often need one-dimensional vectors but in code we want to represent things in richer data formats (e.g. dictionaries of arbitrary dimensional arrays). Prime examples are the parameters of a neural network. Pytrees are not a specific type themselves. For our purposes, any (nested) container of tensors or numbers is a pytree and the pytree operations are defined on it.
The most important operations are:
- `tree_flatten`: Convert a pytree into a list of leaves and an object that stores the tree structure
- `tree_unflatten`: The inverse of `tree_flatten`
- `tree_map`: Apply a function to one or several pytrees
- `tree_reduce`: Apply a reduction over all leaves in a pytree
A full list is [here](https://jax.readthedocs.io/en/latest/jax.tree_util.html)
What is considered a leave depends on the registry of containers, which can be extended by users. For example, if torch tensors are not registered containers, `tree_flatten` would convert a nested dict of tensors into a list of tensors. If torch tensors are registered, it would flatten a nested dict of tensors into a list of numbers.
## pydvl code that could be removed or simplified
- `pydvl.influence.torch.base.DictBilinearForm`
- `pydvl.influence.torch.base.TensorDictOperator`
- `pydvl.influence.torch.util.reshape_vector_to_tensors`
- `pydvl.influence.torch.util.align_structure`
- `pydvl.influence.torch.util.align_with_model`
- `pydvl.influence.torch.util.flatten_dimensions`
## Advantage of using pytree operations
- They are very general and stay the same across project, whereas custom implementations of similar logic are often project specific.
- They are well designed and by composing multiple pytree operations one can achieve a lot
- They are optimized (especially in [optree](https://github.com/metaopt/optree?tab=readme-ov-file#tree-map-nargs))
- They can help us to achieve framework independence. By playing with the container registry, the operations stay the same for pytorch and jax models (even though jax models will typically have a more nested parameter representation)
## Drawbacks of using pytree operations
- We get an additional dependency
- There is a learning curve to pytrees (but it's very much worth it!)
Contributor guide
Research direction
The issue points to specific files in pydvl.influence.torch.base and pydvl.influence.torch.util. Start by reading the documentation for pytree operations in jax or optree. Examine the listed functions to understand their current implementations. The goal is to replace custom flattening/alignment logic with pytree operations, ensuring the same functionality. Test changes by running existing influence function tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- pytorch
- Domain
- backend, tooling
- Issue type
- Refactor
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100