pybaum provides tools to work with pytrees which is a concept burrowed from JAX.

Last update: May 11, 2022

Related tags

Deep Learning pybaum

Overview

pybaum

https://readthedocs.org/projects/pybaum/badge/?version=latest

https://img.shields.io/github/workflow/status/OpenSourceEconomics/pybaum/main/main

Installation

pybaum is available on PyPI and Anaconda.org. Install it with

$ pip install pybaum

# or

$ conda install -c conda-forge pybaum

About

pybaum provides tools to work with pytrees which is a concept burrowed from JAX.

What is a pytree?

In pybaum, we use the term pytree to refer to a tree-like structure built out of container-like Python objects. Classes are considered container-like if they are in the pytree registry, which by default includes lists, tuples, and dicts. That is:

Any object whose type is not in the pytree container registry is considered a leaf pytree.
Any object whose type is in the pytree container registry, and which contains pytrees, is considered a pytree.

For each entry in the pytree container registry, a container-like type is registered with a pair of functions that specify how to convert an instance of the container type to a (children, metadata) pair and how to convert such a pair back to an instance of the container type. Using these functions, pybaum can canonicalize any tree of registered container objects into tuples.

Example pytrees:

[1, "a", object()]  # 3 leaves

(1, (2, 3), ())  # 3 leaves

[1, {"k1": 2, "k2": (3, 4)}, 5]  # 5 leaves

pybaum can be extended to consider other container types as pytrees.

Comments

Fix for new jax version.

Problem

The type of a jax array has changed in jax version 0.4.1 and thus pybaum would not work with this jax version.

Solution

Handle jax arrays explicitly in get_type, as we are currently already doing for namedtuples.

opened by janosg 1
[pre-commit.ci] pre-commit autoupdate
updates:

github.com/pre-commit/pre-commit-hooks: v4.3.0 → v4.4.0

github.com/asottile/reorder_python_imports: v3.1.0 → v3.9.0

github.com/pre-commit/pre-commit-hooks: v4.3.0 → v4.4.0

github.com/psf/black: 22.3.0 → 22.12.0

github.com/PyCQA/flake8: 4.0.1 → 6.0.0

github.com/PyCQA/doc8: 0.11.2 → v1.0.0

github.com/mgedmin/check-manifest: 0.48 → 0.49

github.com/PyCQA/doc8: 0.11.2 → v1.0.0

github.com/asottile/setup-cfg-fmt: v1.20.1 → v2.2.0

github.com/codespell-project/codespell: v2.1.0 → v2.2.2

github.com/asottile/pyupgrade: v2.34.0 → v3.3.1
opened by pre-commit-ci[bot] 1
[pre-commit.ci] pre-commit autoupdate
updates:

github.com/pre-commit/pre-commit-hooks: v4.2.0 → v4.3.0

github.com/pre-commit/pre-commit-hooks: v4.2.0 → v4.3.0

github.com/PyCQA/doc8: 0.11.1 → 0.11.2

github.com/PyCQA/doc8: 0.11.1 → 0.11.2

github.com/asottile/pyupgrade: v2.32.1 → v2.34.0
opened by pre-commit-ci[bot] 1
[pre-commit.ci] pre-commit autoupdate
updates:

github.com/pre-commit/pre-commit-hooks: v4.1.0 → v4.2.0

github.com/asottile/reorder_python_imports: v2.7.1 → v3.1.0

github.com/pre-commit/pre-commit-hooks: v4.1.0 → v4.2.0

github.com/psf/black: 22.1.0 → 22.3.0

github.com/PyCQA/doc8: 0.10.1 → 0.11.1

github.com/mgedmin/check-manifest: 0.47 → 0.48

github.com/PyCQA/doc8: 0.10.1 → 0.11.1

github.com/asottile/setup-cfg-fmt: v1.20.0 → v1.20.1

github.com/asottile/pyupgrade: v2.31.0 → v2.32.1
opened by pre-commit-ci[bot] 1
Align default containers with jax
Currently we have only list, dict and tuple as default pytree containers.

Jax also has:

[x] None

[x] collections.namedtuple and typing.NamedTuple

[x] collections.OrderedDict
opened by janosg 1
[pre-commit.ci] pre-commit autoupdate
updates:

github.com/pre-commit/pre-commit-hooks: v4.0.1 → v4.1.0

github.com/asottile/reorder_python_imports: v2.6.0 → v2.7.1

github.com/pre-commit/pre-commit-hooks: v4.0.1 → v4.1.0

github.com/asottile/blacken-docs: v1.12.0 → v1.12.1

github.com/psf/black: 21.12b0 → 22.1.0
opened by pre-commit-ci[bot] 1
Additional option for flattening pandas DataFrames
Currently the user can decide between two ways of treating DataFrames

Only the entries in a fixed set of columns appears in the flat list of leaves

All entries appear in the flat list of leaves.

In estimagic a third mode would be very practical that automatically switches between the two cases above. If all columns from the specified subset are present, only those columns are considered. If None of the columns is present, all entries are considered.

If this sounds too special, we can also only add it in estimagic.
opened by janosg 1
Add basic documentation
To-Do

[x] Add all functions that can be directly imported from pybaum to an API Documentation

[ ] Write some text about the rationale behind pybaum

[ ] Explain deviations from JAX

[ ] Add examples

enhancement good first issue
opened by janosg 0
Test against JAX
To-Do

[ ] Write sample_tree function that creates a random pytree based on some high level information

[ ] Create random pytrees and test against JAX

[ ] Test current example pytrees against JAX

Note

Need to find a way of not installing JAX on windows in tox.
enhancement good first issue
opened by janosg 0
Improve treedef
Problem

The JAX treedef contains more information. Therefore, their tree_unflatten function does not need an is_leaf argument. The same way we could make the registry argument unnecessary.

This is not just more convenient. It reduces error potential because the unflattening will automatically use the same options as the flattening.

To-Do

[ ] Come up with better Treedef (namedtuple or Class) and implement it

[ ] Adjust interface of tree_unflatten.

enhancement
opened by janosg 1
Improve error checking
To-Do

[ ] Add explicit error handling to tree_equal

[ ] Check for typos in options and types in get_registry and suggest improvements with fuzzywuzzy

[ ] Look for more places where we need better error handling

enhancement
opened by janosg 0

Releases(v0.1.3)

v0.1.3(Dec 16, 2022)
#22 Improves the type check for jax arrays and enables support for Jax 0.4.1 or newer.

Source code(tar.gz)
Source code(zip)
v0.1.2(Jun 15, 2022)
Add JAX array support to pybaum. By @timmens and @janosg in #18 and #20

Source code(tar.gz)
Source code(zip)
v0.1.1(Feb 19, 2022)
What's Changed

Add a section about the API to the docs. by @tobiasraabe in https://github.com/OpenSourceEconomics/pybaum/pull/12

Add default arguments to tree_just_flatten. by @tobiasraabe in https://github.com/OpenSourceEconomics/pybaum/pull/10

Implement tree_yield and tree_just_yield. by @tobiasraabe in https://github.com/OpenSourceEconomics/pybaum/pull/11

Add badges and a short description to the readme. by @tobiasraabe in https://github.com/OpenSourceEconomics/pybaum/pull/13

Full Changelog: https://github.com/OpenSourceEconomics/pybaum/compare/v0.1.0...v0.1.1
Source code(tar.gz)
Source code(zip)
v0.1.0(Jan 28, 2022)

Source code(tar.gz)
Source code(zip)

Owner

Open Source Economics

A platform for transdisciplinary collaboration in the development, analysis, and application of computational economic models.

GitHub Repository https://pybaum.readthedocs.io

This repository contains the source codes for the paper AtlasNet V2 - Learning Elementary Structures.

AtlasNet V2 - Learning Elementary Structures This work was build upon Thibault Groueix's AtlasNet and 3D-CODED projects. (you might want to have a loo

123 Nov 11, 2022

Convert BART models to ONNX with quantization. 3X reduction in size, and upto 3X boost in inference speed

fast-Bart Reduction of BART model size by 3X, and boost in inference speed up to 3X BART implementation of the fastT5 library (https://github.com/Ki6a

19 Dec 09, 2022

Exploring Versatile Prior for Human Motion via Motion Frequency Guidance (3DV2021)

Exploring Versatile Prior for Human Motion via Motion Frequency Guidance [Video Demo] [Paper] Installation Requirements Python 3.6 PyTorch 1.1.0 Pleas

19 Oct 28, 2022

Decoding the Protein-ligand Interactions Using Parallel Graph Neural Networks

Decoding the Protein-ligand Interactions Using Parallel Graph Neural Networks Requirements python 0.10+ rdkit 2020.03.3.0 biopython 1.78 openbabel 2.4

3 Nov 23, 2022

🥈78th place in Riiid Solution🥈

Riiid Answer Correctness Prediction Introduction This repository is the code that placed 78th in Riiid Answer Correctness Prediction competition. Requ

14 Apr 26, 2022

Simple Pose: Rethinking and Improving a Bottom-up Approach for Multi-Person Pose Estimation

SimplePose Code and pre-trained models for our paper, “Simple Pose: Rethinking and Improving a Bottom-up Approach for Multi-Person Pose Estimation”, a

256 Dec 24, 2022

Planning from Pixels in Environments with Combinatorially Hard Search Spaces -- NeurIPS 2021

PPGS: Planning from Pixels in Environments with Combinatorially Hard Search Spaces Environment Setup We recommend pipenv for creating and managing vir

11 Jun 26, 2022

Disagreement-Regularized Imitation Learning

Due to a normalization bug the expert trajectories have lower performance than the rl_baseline_zoo reported experts. Please see the following link in

25 Apr 28, 2022

Point Cloud Registration Network

PCRNet: Point Cloud Registration Network using PointNet Encoding Source Code Author: Vinit Sarode and Xueqian Li Paper | Website | Video | Pytorch Imp

59 Nov 19, 2022

custom pytorch implementation of MoCo v3

MoCov3-pytorch custom implementation of MoCov3 [arxiv]. I made minor modifications based on the official MoCo repository [github]. No ViT part code an

39 Nov 14, 2022

A light and fast one class detection framework for edge devices. We provide face detector, head detector, pedestrian detector, vehicle detector......

A Light and Fast Face Detector for Edge Devices Big News: LFD, which is a big update of LFFD, now is released (2021.03.09). It is strongly recommended

1.3k Dec 25, 2022

pybaum provides tools to work with pytrees which is a concept burrowed from JAX.

Related tags

Overview

pybaum

Installation

About

Comments

Problem

Solution

To-Do

To-Do

Note

Problem

To-Do

To-Do

Releases(v0.1.3)

v0.1.3(Dec 16, 2022)

v0.1.2(Jun 15, 2022)

v0.1.1(Feb 19, 2022)

What's Changed

v0.1.0(Jan 28, 2022)

Owner

Open Source Economics

This repository contains the source codes for the paper AtlasNet V2 - Learning Elementary Structures.

Convert BART models to ONNX with quantization. 3X reduction in size, and upto 3X boost in inference speed

Exploring Versatile Prior for Human Motion via Motion Frequency Guidance (3DV2021)

Decoding the Protein-ligand Interactions Using Parallel Graph Neural Networks

🥈78th place in Riiid Solution🥈

Simple Pose: Rethinking and Improving a Bottom-up Approach for Multi-Person Pose Estimation

Planning from Pixels in Environments with Combinatorially Hard Search Spaces -- NeurIPS 2021

Disagreement-Regularized Imitation Learning

Point Cloud Registration Network

custom pytorch implementation of MoCo v3

A light and fast one class detection framework for edge devices. We provide face detector, head detector, pedestrian detector, vehicle detector......

DEEPAGÉ: Answering Questions in Portuguese about the Brazilian Environment

[ICCV 2021 Oral] Mining Latent Classes for Few-shot Segmentation

CoCosNet v2: Full-Resolution Correspondence Learning for Image Translation

Res2Net for Instance segmentation and Object detection using MaskRCNN

Puzzle-CAM: Improved localization via matching partial and full features.

The PyTorch implementation of Directed Graph Contrastive Learning (DiGCL), NeurIPS-2021

Time series annotation library.

[CVPR 2022] Unsupervised Image-to-Image Translation with Generative Prior

Management Dashboard for Torchserve