Skip to content

Commit

Permalink
replace pyrsistent.pmap with immutabledict (#248)
Browse files Browse the repository at this point in the history
* replace pyrsistent.pmap with immutabledict

* change type annotation to 'Mapping'
  • Loading branch information
matthiasdiener authored Oct 9, 2023
1 parent a62e26d commit a5fe5b2
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 9 deletions.
14 changes: 6 additions & 8 deletions arraycontext/impl/pytato/compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
from typing import Any, Callable, Dict, FrozenSet, Mapping, Tuple, Type

import numpy as np
from pyrsistent import PMap, pmap
from immutabledict import immutabledict

import pytato as pt
from pytools import ProcessLogger
Expand Down Expand Up @@ -131,11 +131,9 @@ def _rec_str(key: Any) -> str:

def _get_arg_id_to_arg_and_arg_id_to_descr(args: Tuple[Any, ...],
kwargs: Mapping[str, Any]
) -> "Tuple[PMap[Tuple[Any, ...],\
Any],\
PMap[Tuple[Any, ...],\
AbstractInputDescriptor]\
]":
) -> \
Tuple[Mapping[Tuple[Any, ...], Any],
Mapping[Tuple[Any, ...], AbstractInputDescriptor]]:
"""
Helper for :meth:`BaseLazilyCompilingFunctionCaller.__call__`. Extracts
mappings from argument id to argument values and from argument id to
Expand Down Expand Up @@ -171,7 +169,7 @@ def id_collector(keys, ary):
" either a scalar, pt.Array or an array container. Got"
f" '{arg}'.")

return pmap(arg_id_to_arg), pmap(arg_id_to_descr)
return immutabledict(arg_id_to_arg), immutabledict(arg_id_to_descr)


def _to_input_for_compiled(ary: ArrayT, actx: PytatoPyOpenCLArrayContext):
Expand Down Expand Up @@ -259,7 +257,7 @@ class BaseLazilyCompilingFunctionCaller:

actx: _BasePytatoArrayContext
f: Callable[..., Any]
program_cache: Dict["PMap[Tuple[Any, ...], AbstractInputDescriptor]",
program_cache: Dict[Mapping[Tuple[Any, ...], AbstractInputDescriptor],
"CompiledFunction"] = field(default_factory=lambda: {})

# {{{ abstract interface
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ def main():

# https://github.com/inducer/arraycontext/pull/147
"pytools>=2022.1.3",

"immutabledict",
"loopy>=2019.1",
],
extras_require={
Expand Down

0 comments on commit a5fe5b2

Please sign in to comment.