Matt Smith authored 3 days ago and Andreas Klöckner committed 1 day ago

add comment explaining why CachedMapper.get_cache_key and get_function_definition_cache_key are not defined for general extra args/kwargs

Matt Smith authored 1 month ago and Andreas Klöckner committed 1 day ago

disable default implementation of get_cache_key and get_function_definition_cache_key for extra args case

ambiguous due to the fact that any arg can be specified with/without keyword

* refactor deduplicate_data_wrappers to avoid dependence on erroneous super().rec usage in CachedMapAndCopyMapper

Here is a sketch of what happens with super().rec vs Mapper.rec for the previous implementation of
deduplicate_data_wrappers. Suppose we have 2 data wrappers a and b with the same data pointer.

With super().rec:

1) map_fn maps a to itself, then mapper copies a to a'; mapper caches a -> a' (twice, once in
   super().rec and then again in rec),
2) map_fn maps b to a, then mapper maps (via cache in super().rec call) a to a'; mapper caches
   b -> a'.

=> Only a' in output DAG.

With Mapper.rec:

1) map_fn maps a to itself, then mapper copies a to a'; caches a -> a',
2) map_fn maps b to a, then mapper copies a to a''; caches b -> a''.

=> Both a' and a'' in output DAG.

* call Mapper.rec instead of super().rec to avoid double caching

* call Mapper.rec from CachedMapper too just to avoid copy/paste errors

* add assertion to check for double caching

* add comment explaining use of Mapper.rec

Andreas Klöckner authored 1 month ago and Andreas Klöckner committed 1 month ago

- Better handle 1-long axes in old and new shapes
- Avoid generating modulo expressions for direct pass-through

* make CombineMapper inherit from CachedMapper

* move axis tag attaching code into a separate method in AxisTagAttacher

* use a class for CachedMapper caches instead of using a dict directly

* drop P from mapper caches, since it's not being used

* add comment about type annotations in CachedMapperCache

* avoid traversing functions multiple times

* condense function cache setup

* pass function cache to CombineMapper.__init__ too

* use Never instead of None as the function result type for mappers that don't support functions

* remove multiple inheritance for TransformMapperWithExtraArgs

doesn't appear to be needed

* add _verify_is_array to avoid the need for rec_ary

the latter inflates recursion depth

* remove map_foreign from Mapper

* use Never as FunctionResultT for NumpyCodegenMapper and FancyDotWriter too

* make PlaceholderSubstitutor explicit about not supporting functions

* tweak types in rec/rec_function_definition

* use P.args/P.kwargs in CachedWalkMapper

* use a more specific type than Any for cache key in CachedWalkMapper

* remove some more Anys

* do not pickle cached hash of arrays

* no clue why this works

* remove unneeded changes

* add an explanation

* add simple test

* add another test

* copy getstate/setstate implementation from dataclasses

* typo?

* DictOfNamedArrays: better repr, keys

* add tags to repr

* make keys() abstract

* add to doc

* fix lint

* remove invalid requirements.txt pytools version specification

---------

Co-authored-by: Andreas Klöckner <inform@tiker.net>

* change attrs to dataclasses

* post-merge fix

* more fixes

* switch to array_dataclass

* use _augment_array_dataclass

* restore attribute validation

* mypy?

* Disable pylint invalid-field-call

---------

Co-authored-by: Andreas Kloeckner <inform@tiker.net>