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

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

Deep High-Resolution Representation Learning for Human Pose Estimation

Fast mesh denoising with data driven normal filtering using deep variational autoencoders

Shape-aware Semi-supervised 3D Semantic Segmentation for Medical Images

A comprehensive list of published machine learning applications to cosmology

Ganilla - Official Pytorch implementation of GANILLA

50-days-of-Statistics-for-Data-Science - This repository consist of a 50-day program

DGN pymarl - Implementation of DGN on Pymarl, which could be trained by VDN or QMIX

Official implementation of the method ContIG, for self-supervised learning from medical imaging with genomics

Hyper-parameter optimization for sklearn

MG-GCN: Scalable Multi-GPU GCN Training Framework

CCPD: a diverse and well-annotated dataset for license plate detection and recognition

UI2I via StyleGAN2 - Unsupervised image-to-image translation method via pre-trained StyleGAN2 network

Run Keras models in the browser, with GPU support using WebGL

My implementation of Fully Convolutional Neural Networks in Keras

QT Py Media Knob using rotary encoder & neopixel ring

Semantically Contrastive Learning for Low-light Image Enhancement

This program was designed to detect whether someone is wearing a facemask through a live video stream.

State of the art Semantic Sentence Embeddings

UnpNet - Rethinking 3-D LiDAR Point Cloud Segmentation(IEEE TNNLS)

AdvStyle - Official PyTorch Implementation