Visualize, create, and operate on pytrees in the most intuitive way possible.
APACHE-2.0 License
Bot releases are hidden (Show)
Published by ASEM000 over 1 year ago
alias
in Field
by @ASEM000 in https://github.com/ASEM000/PyTreeClass/pull/50
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.2.5...v0.2.6
Published by ASEM000 over 1 year ago
fix unordered keys bug in registering instance variables
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.2.4...v0.2.5
Published by ASEM000 over 1 year ago
...
in .at[]
to work with is_leaf
import pytreeclass as pytc
@pytc.treeclass
class Tree:
a:int = 1
tree = Tree()
tree.at[...].set(is_leaf=...) # is leaf works with ...
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.2.3...v0.2.4
Published by ASEM000 over 1 year ago
tree_repr_with_trace
in #47 (useful for debugging)a:list=[]
but not for a:list=pytc.field(default=[])
default_factory
-> factory
( recognized by dataclass_transform
and attrs
)self
lru_cache
to improve treeclass
perf.Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.2.2...v0.2.3
Published by ASEM000 over 1 year ago
import pytreeclass as pytc
tree = [1, [2, [3]]]
traces, _ = zip(*pytc.tree_leaves_with_trace(tree))
print(pytc.tree_repr([trace[3] for trace in traces]))
# [
# ({id:4843302976}, {id:4307173616}),
# ({id:4843302976}, {id:4843303936}, {id:4307173648}),
# ({id:4843302976}, {id:4843303936}, {id:4843297408}, {id:4307173680})
# ]
is_tree_equal(*trees)
now accepts arbitrary number of pytrees.FrozenWrapper
is private.Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.2.1...v0.2.2
Published by ASEM000 over 1 year ago
dataclass_transform
for better typing/autocompleteFull Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.2.0...v0.2.1
Published by ASEM000 over 1 year ago
Mostly rewritten from scratch.
Drop the standard library dataclasses.dataclass
for faulty logic treating jax.numpy as mutable object from python 3.11 see: #33
Implements a simpler dataclass
-like version that additionally implements:
pos_only
to mark the argument as positional onlycallbacks
to apply a set of functions (e.g., validators/converters) on input see: #37Drop metadata-based approach for marking fields static. move to a transparent wrapper approach for faster, simpler, and flexible marking of leaves static see here
tree_summary
, tree_diagram
, tree_str
, and tree_repr
can now work on any pytree
(including registered pytrees) out of the box see
Remove tree.summary()
in favor of pytc.tree_summary(model)
Remove tree.tree_diagram()
in favor of pytc.tree_diagram(model)
Remove tree.tree_box()
Add tree_{repr,str,diagram,summary,mermaid}(..., depth)
to control cutoff depth condition for visualization.
Make math operations on leaves opt-in feature via treeckass(... , leafwise=True)
Drop filtering by metadata
, types
, and name
. Instead, use tree_map_with_trace
to access leaves' name,type,index, and metadata at each level along their path. see #35
Add bcmap
: a function transformation for automatic broadcasting. see readme More section.
bcmap
with treeclass(... , leafwise=True)
achieves a data model that enables applying numpy
functions like numpy.where
on arbitrary PyTrees without the need for tree_map
.Faster flatten/unflatten see benchmarks comparison with flax/equinox.
Internally, unifying indexing and viz tools by implementing custom jax-like registry to record names, types, indices, and metadata rules for common python data structures and treeclass
wrapped class. You can check the readme More section for registering custom classes to work fully with the viz/indexing tools similar to Jax registration process.
Published by ASEM000 over 1 year ago
Initial beta release for v0.2
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.13...0.2.0b13
Published by ASEM000 about 2 years ago
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.12...v0.1.13
Published by ASEM000 about 2 years ago
__delattr__
pytc.is_frozen_field
-> pytc.is_field_frozen
pytc.is_nondiff_field
-> pytc.is_field_nondiff
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.10...v0.1.11
Published by ASEM000 about 2 years ago
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.9...v0.1.10
tree_mermaid
placed on tree arrows.pytc.field
wrapping dataclasses.field
with extra nondiff, frozen
argumentspytc.fields
similar to dataclasses.fields
which returns dataclasses
fields and the undeclared fieldspytc.is_treeclass_{frozen, leaf_bool, leaf, equal}
tree_diagram
and tree_summary
if at least one treeclass
item exist ..at[].set(...)
Published by ASEM000 about 2 years ago
array_as_leaves
option in .at
methodsisinstance
for speed ( notable speed improvements in all treeviz methods )tree_mermaid
.flatten/unflatten
had to be declared as a dataclass
field. Now by default, all non-treeclass
instance variables are captured as static(i.e. part of treedef
) unless declared in the dataclass
fields.import pytreeclass as pytc
import jax.tree_util as jtu
@pytc.treeclass
class Test:
def __init__(self):
self.a=1
t = Test()
tt = jtu.tree_unflatten(*jtu.tree_flatten(t)[::-1])
# before change it will raise an error as
# `a` was not captured since it was not declared as dataclass field
tt.a # -> Error
# After change, it will be part of the treedef (i.e. static/nondiff field)
tt.a # 1
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.8...v0.1.9
Published by ASEM000 about 2 years ago
.at
methods ( up to 40% improvement in speed )Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.7...v0.1.8
Published by ASEM000 about 2 years ago
treeclass
wrapped class is treated as a parameter i.e. field_only=False
. This will simplify reasoning about parameter management on the user end.__undeclared_fields__
with MappingProxyType
to make it immutable.Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/0.1.6...v0.1.7
Published by ASEM000 about 2 years ago
one addition and one change in this release,
filter_nondiff
, unfilter_nondiff
are added. see Readme for example usage.tree_freeze
now marks all fields static instead of caching them. this implementation is slower but it unifies the code and remove the distinction between module and fields when marking static.Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.5...0.1.6
Published by ASEM000 about 2 years ago
mask_util
.__str__
and __repr__
.Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/0.1.4...v0.1.5
Published by ASEM000 about 2 years ago
{key:val} == tree
maps {key:val} == field.metadata
for each field item. Previously the operation was {key:val} in field.metadata
.at
aux methods namely add,multiply,power,divide,max,min,reduce_sum,reduce_product,reduce_max,reduce_min
.param
and Container.
these features will be moved to serket
in the future._treeBase
methods : __pytree_fields__
, __pytree_structure__
, frozen
property, is replaced with _tree_fields
, _tree_structure
, _is_treeclass_frozen
functions in tree_util
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/0.1.3...0.1.4
Published by ASEM000 about 2 years ago
.at[].freeze
to pytreeclass.tree_util.tree_freeze
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.2...0.1.3
Published by ASEM000 about 2 years ago
Full Changelog: https://github.com/ASEM000/PyTreeClass/compare/v0.1.1...v0.1.2
Published by ASEM000 about 2 years ago
PyTreeClass
first minor release.at
for string/boolean argument
model.at["a"].at["b"].{get,set,apply,reduce,...} <-> model.a.b
model.at[model>0].at[model == "linear"].{get,set,apply,reduce,...}
<-> model.at[(model>0) & (model=="linear")].{get,set,apply,reduce,...}
black
-like str/reprstatic_value
filter_nondiff
to make non-inexact values static. and , unfilter_nondiff
to revert that operation. #21register_node
to param
_pytree_map
to operate on dataclass
field and value