Intelegentny_Pszczelarz/.venv/Lib/site-packages/jax/_src/linear_util.py

399 lines
13 KiB
Python
Raw Normal View History

2023-06-19 00:49:18 +02:00
# Copyright 2018 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Utilities for defining functions composed with transformations.
For example,
from jax._src import linear_util as lu
wf = lu.wrap_init(f) # Produce a WrappedFun for applying transformations on `f`
A `WrappedFun` object represents a function `f`, together with a sequence of
nested transformations that are to be applied to the positional and keyword
arguments at call time and function return values at return time.
A transformation can take some static positional arguments that are given
at the wrapping time, and may also return some auxiliary output:
wf, aux_out_thunk = trans1(wf, static_arg)
We can call the transformed function. First, the transformation is applied
to the dynamic args and keyword args to produce new dynamic and keyword args.
Then the underlying function is called and the transformation is applied to
the results.
If there are multiple transformations, they form a stack. The arguments are
transformed first with the last applied transformation; the results are
transformed first with the first applied transformation.
res = wf.call_wrapped(dynamic_args, kwargs)
# Now `aux_out_thunk()` is the auxiliary output.
A transformation is written as a generator function that takes zero or more
static positional arguments (given when the transformation is instantiated),
along with positional and keyword arguments to be transformed.
The generator will yield twice:
@lu.transformation_with_aux
def trans1(static_arg, *dynamic_args, **kwargs):
...
# First yield: pair of transformed (args, kwargs). Get back the results.
results = yield (new_dynamic_args, new_kwargs)
...
# Second yield: pair of (transformed results, and auxiliary output)
yield new_results, auxiliary_output
`WrappedFun` objects explicitly represent the set of transformations so that
they can be used as dictionary keys for memoization. `WrappedFun` objects
compare as equal only if they compute the same function. The static and the
dynamic positional arguments for the generators, and also the auxiliary output
data must be immutable, because it will be stored in function memoization tables.
"""
from __future__ import annotations
from functools import partial
import operator
from typing import Any, Tuple, Callable, Optional, NamedTuple
import weakref
from jax._src.tree_util import tree_map
from jax._src.config import config
from jax._src import core
from jax._src import traceback_util
from jax._src.util import curry
traceback_util.register_exclusion(__file__)
class StoreException(Exception): pass
class EmptyStoreValue: pass
_EMPTY_STORE_VALUE = EmptyStoreValue()
class Store:
"""Storage for a value, with checks for overwriting or reading empty store."""
__slots__ = ("_val",)
def __init__(self):
self._val = _EMPTY_STORE_VALUE
def store(self, val):
if self._val is not _EMPTY_STORE_VALUE:
raise StoreException("Store occupied")
self._val = val
def reset(self):
# This should only be called in exceptional circumstances (e.g. debugging).
self._val = _EMPTY_STORE_VALUE
@property
def val(self):
if not self:
raise StoreException("Store empty")
return self._val
def __nonzero__(self):
return self._val is not _EMPTY_STORE_VALUE
__bool__ = __nonzero__
class EqualStore:
__slots__ = ('_store',)
def __init__(self):
self._store = Store()
val = property(operator.attrgetter('_store.val'))
def store(self, val):
try:
self._store.store(val)
except StoreException as e:
try:
okay = bool(self._store._val == val)
except:
raise e from None
else:
if not okay:
raise StoreException("Store occupied with not-equal value") from None
class WrappedFun:
"""Represents a function `f` to which `transforms` are to be applied.
Args:
f: the function to be transformed.
transforms: a list of `(gen, gen_static_args)` tuples representing
transformations to apply to `f.` Here `gen` is a generator function and
`gen_static_args` is a tuple of static arguments for the generator. See
description at the start of this module for the expected behavior of the
generator.
stores: a list of out_store for the auxiliary output of the `transforms`.
params: extra parameters to pass as keyword arguments to `f`, along with the
transformed keyword arguments.
"""
__slots__ = ("f", "transforms", "stores", "params", "in_type", "debug_info")
def __init__(self, f, transforms, stores, params, in_type, debug_info):
self.f = f
self.transforms = transforms
self.stores = stores
self.params = params
self.in_type = in_type
self.debug_info = debug_info
@property
def __name__(self):
return getattr(self.f, '__name__', '<unnamed wrapped function>')
def wrap(self, gen, gen_static_args, out_store) -> WrappedFun:
"""Add another transform and its store."""
return WrappedFun(self.f, ((gen, gen_static_args),) + self.transforms,
(out_store,) + self.stores, self.params, None, None)
def populate_stores(self, stores):
"""Copy the values from the `stores` into `self.stores`."""
for self_store, other_store in zip(self.stores, stores):
if self_store is not None:
self_store.store(other_store.val)
def call_wrapped(self, *args, **kwargs):
"""Calls the underlying function, applying the transforms.
The positional `args` and keyword `kwargs` are passed to the first
transformation generator.
"""
stack = []
for (gen, gen_static_args), out_store in zip(self.transforms, self.stores):
gen = gen(*(gen_static_args + tuple(args)), **kwargs)
args, kwargs = next(gen)
stack.append((gen, out_store))
gen = gen_static_args = out_store = None
try:
ans = self.f(*args, **dict(self.params, **kwargs))
except:
# Some transformations yield from inside context managers, so we have to
# interrupt them before reraising the exception. Otherwise they will only
# get garbage-collected at some later time, running their cleanup tasks
# only after this exception is handled, which can corrupt the global
# state.
while stack:
stack.pop()[0].close()
raise
args = kwargs = None
while stack:
gen, out_store = stack.pop()
try:
ans = gen.send(ans)
except:
# As above does for the first half of the transformation, exceptions
# raised in the second half of the transformation also require us to
# clean up references here.
while stack:
stack.pop()[0].close()
raise
if out_store is not None:
ans, side = ans
out_store.store(side)
return ans
def __repr__(self):
def transform_to_str(x):
i, (gen, args) = x
return f"{i} : {fun_name(gen)} {fun_name(args)}"
transformation_stack = map(transform_to_str, enumerate(self.transforms))
return "Wrapped function:\n" + '\n'.join(transformation_stack) + '\nCore: ' + fun_name(self.f) + '\n'
def __hash__(self):
return hash((self.f, self.transforms, self.params, self.in_type,
self.debug_info))
def __eq__(self, other):
return (self.f == other.f and self.transforms == other.transforms and
self.params == other.params and self.in_type == other.in_type and
self.debug_info == other.debug_info)
@curry
def transformation(gen, fun: WrappedFun, *gen_static_args) -> WrappedFun:
"""Adds one more transformation to a WrappedFun.
Args:
gen: the transformation generator function
fun: a WrappedFun on which to apply the transformation
gen_static_args: static args for the generator function
"""
return fun.wrap(gen, gen_static_args, None)
@curry
def transformation_with_aux(gen, fun: WrappedFun, *gen_static_args,
use_eq_store=False) -> Tuple[WrappedFun, Any]:
"""Adds one more transformation with auxiliary output to a WrappedFun."""
out_store = Store() if not use_eq_store else EqualStore()
out_thunk = lambda: out_store.val
return fun.wrap(gen, gen_static_args, out_store), out_thunk
def fun_name(f):
try:
return f.__name__
except:
return str(f)
def wrap_init(f, params=None) -> WrappedFun:
"""Wraps function `f` as a `WrappedFun`, suitable for transformation."""
params = () if params is None else tuple(sorted(params.items()))
return WrappedFun(f, (), (), params, None, None)
def annotate(f: WrappedFun, in_type: Optional[core.InputType]) -> WrappedFun:
assert f.in_type is None
if in_type is None:
return f
_check_input_type(in_type)
return WrappedFun(f.f, f.transforms, f.stores, f.params, in_type, f.debug_info)
def _check_input_type(in_type: core.InputType) -> None:
# Check that in_type is syntactically well-formed
assert type(in_type) is tuple and all(type(e) is tuple for e in in_type)
assert all(isinstance(a, core.AbstractValue) and type(b) is bool
and not isinstance(a, core.ConcreteArray) for a, b in in_type)
def valid_size(d) -> bool:
if isinstance(d, core.DBIdx) and type(d.val) is int and d.val >= 0:
return True
return (isinstance(d, (int, core.DBIdx, core.DArray)) and
(not isinstance(d, core.DArray) or type(d) is core.bint and not d.shape))
assert all(valid_size(d) for a, _ in in_type if type(a) is core.DShapedArray
for d in a.shape)
# Check that all DBIdx point to positions to the left of the input on which
# they appear.
assert all(d.val < i for i, (aval, _) in enumerate(in_type)
if isinstance(aval, core.DShapedArray) for d in aval.shape
if isinstance(d, core.DBIdx))
# Check that all implicit arguments have at least one DBIdx pointing to them.
provided = [e for _, e in in_type]
for aval, _ in in_type:
if type(aval) is core.DShapedArray:
for d in aval.shape:
if isinstance(d, core.DBIdx):
provided[d.val] = True
assert all(provided)
class TracingDebugInfo(NamedTuple):
# Packages up trace/staging-time debug info about a func and its parameters,
# formed just before staging to a jaxpr and read in trace-time error messages.
# TODO(mattjj): delete partial_eval.DebugInfo, replace all uses with this cls
traced_for: str # e.g. 'jit', 'scan', etc
func_src_info: str # e.g. f'{fun.__name__} at {filename}:{lineno}'
arg_names: Tuple[str, ...] # e.g. ('args[0]', ... )
result_paths: Optional[Callable[[], Tuple[str, ...]]]
def add_debug_info(f: WrappedFun, debug_info: Optional[TracingDebugInfo]
) -> WrappedFun:
"""Produce a new WrappedFun with debug_info attached."""
assert f.debug_info is None
if debug_info is None:
return f
return WrappedFun(f.f, f.transforms, f.stores, f.params, f.in_type, debug_info)
def cache(call: Callable):
"""Memoization decorator for functions taking a WrappedFun as first argument.
Args:
call: a Python callable that takes a WrappedFun as its first argument. The
underlying transforms and params on the WrappedFun are used as part of the
memoization cache key.
Returns:
A memoized version of ``call``.
"""
fun_caches: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
def memoized_fun(fun: WrappedFun, *args):
cache = fun_caches.setdefault(fun.f, {})
if config.jax_check_tracer_leaks:
key = (_copy_main_traces(fun.transforms), fun.params, fun.in_type, args,
config.x64_enabled, config.jax_default_device,
config._trace_context())
else:
key = (fun.transforms, fun.params, fun.in_type, args, config.x64_enabled,
config.jax_default_device, config._trace_context())
result = cache.get(key, None)
if result is not None:
ans, stores = result
fun.populate_stores(stores)
else:
ans = call(fun, *args)
cache[key] = (ans, fun.stores)
return ans
def _evict_function(f):
fun_caches.pop(f, None)
memoized_fun.cache_clear = fun_caches.clear # type: ignore
memoized_fun.evict_function = _evict_function # type: ignore
cache_clearing_funs.add(memoized_fun.cache_clear)
return memoized_fun
cache_clearing_funs = weakref.WeakSet() # type: ignore
def clear_all_caches():
global cache_clearing_funs
for clear in cache_clearing_funs:
clear()
@partial(partial, tree_map)
def _copy_main_traces(x):
if isinstance(x, core.MainTrace):
return core.MainTrace(x.level, x.trace_type, **x.payload)
else:
return x
@transformation
def hashable_partial(*args):
yield (yield args, {})
def merge_linear_aux(aux1, aux2):
try:
out1 = aux1()
except StoreException:
# store 1 was not occupied, so store 2 better be
try:
out2 = aux2()
except StoreException:
raise StoreException("neither store occupied") from None
else:
return False, out2
else:
# store 1 was occupied, so let's check store 2 is not occupied
try:
out2 = aux2()
except StoreException:
return True, out1
else:
raise StoreException("both stores occupied")