Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This is a big update.
Exciting new features!
Added
filter_vmap
.filter_vmap
,filter_pmap
#65.)Added
filter_pmap
.filter_vmap
,filter_pmap
#65.)Added pooling layers:
eqx.nn.Pool
eqx.nn.AvgPool1d
eqx.nn.AvgPool2d
eqx.nn.AvgPool3d
eqx.nn.MaxPool1d
eqx.nn.MaxPool2d
eqx.nn.MaxPool3d
Added
tree_serialise_leaves
andtree_deserialise_leaves
.Added
tree_inference
, as a convenience for toggling all inference flags through a model.Refactoring for nicer APIs
filter_{jit,grad,value_and_grad}
now have an easier-to-use API for specifying which arguments have what behaviour.(args, kwargs)
as a single PyTree, then you can specify adefault
,args
,kwargs
separately. In particular this avoids doing messy stuff likefilter_spec=((...), {})
when you had no kwargs.args
andkwargs
against their runtime values. Both the runtime values, and the filter specification, are matched up against the function signature.e.g. you can do
filter_jit(lambda x: x, kwargs=dict(x=True))(1)
, using a keyword argument at JIT-time and a positional argument at call time.filter_jit(fun)
andfilter_jit(default=...)(fun)
will work.tree_at
can now replace subtrees, and not just leaves.tree_at
. #47.)filter
,partition
now support anis_leaf
argument.is_leaf
argument tofilter
andpartition
#68.)Miscellaneous
filter_jit(filter_grad(fun))
twice will no longer lead to unnecessary recompilation: the secondfilter_grad(fun)
instance will be a PyTree that looks like the firstfilter_grad(fun)
instance, and thus we won't get any recompilation.