aai-institute / aai-institute/pyDVL
Use `pytree` operations to simplify code in influence
- Langage dominant
- Python
- Étoiles
- 146
- Forks
- 10
- Métriques de merge des PR
- Aucune PR mergée en 30 j
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!)
Guide de contribution
Ouvrir le guide de contribution
Évaluation
Cette issue n'a pas encore été évaluée.