提交 07832d35 authored 作者: Brandon T. Willard's avatar Brandon T. Willard 提交者: Brandon T. Willard

Make VM an abstract class and fix its docstrings

上级 83e5f3ef
...@@ -10,6 +10,7 @@ import platform ...@@ -10,6 +10,7 @@ import platform
import sys import sys
import time import time
import warnings import warnings
from abc import ABC, abstractmethod
from collections import defaultdict from collections import defaultdict
from aesara.configdefaults import config from aesara.configdefaults import config
...@@ -121,17 +122,18 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re, depend ...@@ -121,17 +122,18 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re, depend
return reallocated_info return reallocated_info
class VM: class VM(ABC):
""" r"""An abstract class for evaluating Aesara programs.
A VM object's __call__ method evaluates an Aesara program.
The `VM.__call__` method evaluates an Aesara program.
The Stack should be considered the reference VM/Linker implementation. `Stack` should be considered the reference `VM`/`Linker` implementation.
It can correctly evaluate all graphs and is the easiest to read. The CVM It can correctly evaluate all graphs and is the easiest to read. The `CVM`
is a port of Stack, and should have the same behavior, but run faster. is a port of `Stack` and should have the same behavior, but run faster.
The CVM's code is harder to read though. The `CVM`'s code is harder to read though.
The other python VMs are maybe not necessary anymore, and don't take The other python `VM`\s are perhaps not necessary anymore, and don't take
advantage of lazy computation, though they still produce the correct advantage of lazy computation, although they still produce the correct
output for lazy nodes. output for lazy nodes.
Parameters Parameters
...@@ -148,18 +150,18 @@ class VM: ...@@ -148,18 +150,18 @@ class VM:
Attributes Attributes
---------- ----------
call_counts call_counts
List of integers, one for each thunk. call_count[i] is the number of List of integers, one for each thunk. ``call_count[i]`` is the number
times thunks[i] was called in the course of computations performed by of times ``thunks[i]`` was called in the course of computations
call_with_timers(). performed by `call_with_timers`.
call_times call_times
List of floats, one for each thunk. call_times[i] is the amount of List of floats, one for each thunk. ``call_times[i]`` is the amount of
runtime spent on thunks[i] in the course of computations performed by runtime spent on ``thunks[i]`` in the course of computations performed
call_with_timers(). by `call_with_timers`.
need_update_inputs : bool need_update_inputs : bool
True indicates that Function.__call__ must implement the feedback from ``True`` indicates that `Function.__call__` must implement the feedback
output storage to input storage. False means it *must not* repeat that from output storage to input storage. ``False`` means it *must not*
feedback. repeat that feedback.
""" """
...@@ -181,31 +183,24 @@ class VM: ...@@ -181,31 +183,24 @@ class VM:
# defaults to 0 (aka False). # defaults to 0 (aka False).
self.need_update_inputs = True self.need_update_inputs = True
@abstractmethod
def __call__(self): def __call__(self):
""" r"""Run the virtual machine.
Run the machine.
Postcondition - all output variables have been computed. VMs vary in
what exactly this means and how it is done.
After this is executed, all the output variables will have been
computed. `VM`\s may vary regarding what exactly this means and how it
is done.
""" """
raise NotImplementedError()
def clear_storage(self): def clear_storage(self):
""" """Free any internal references to temporary variables.
Free any internal references to temporary variables.
Free internal variables and outputs. Essentially, free as much memory
as possible without intefering with the ability to evaluate subsequent
calls.
Essentially, free as much memory as possible without interfering with
the ability to evaluate subsequent calls.
""" """
raise NotImplementedError()
def update_profile(self, profile): def update_profile(self, profile):
""" """Update a profile object."""
Accumulate into the profile object
"""
for node, thunk, t, c in zip( for node, thunk, t, c in zip(
self.nodes, self.thunks, self.call_times, self.call_counts self.nodes, self.thunks, self.call_times, self.call_counts
): ):
...@@ -244,7 +239,6 @@ class Loop(VM): ...@@ -244,7 +239,6 @@ class Loop(VM):
""" """
# Some other part of Aesara query that information
allow_gc = False allow_gc = False
def __call__(self): def __call__(self):
...@@ -271,10 +265,9 @@ class Loop(VM): ...@@ -271,10 +265,9 @@ class Loop(VM):
class LoopGC(VM): class LoopGC(VM):
""" """Unconditional start-to-finish program execution in Python.
Unconditional start-to-finish program execution in Python.
Garbage collection is possible on intermediate results.
Garbage collection is possible on intermediate results.
""" """
def __init__(self, fgraph, nodes, thunks, pre_call_clear, post_thunk_clear): def __init__(self, fgraph, nodes, thunks, pre_call_clear, post_thunk_clear):
...@@ -321,28 +314,30 @@ class LoopGC(VM): ...@@ -321,28 +314,30 @@ class LoopGC(VM):
class Stack(VM): class Stack(VM):
""" """Finish-to-start evaluation order of thunks.
Finish-to-start evaluation order of thunks.
This supports lazy evaluation of subtrees and partial computations of
graphs when only some inputs have changed.
This supports lazy evaluation of subtrees and partial At a pseudo-code level, the basic idea is as follows:
computations of graphs when only some inputs have changed.
At a pseudo-code level, the basic idea is the following: .. code-block:: python
def recursively_evaluate(var): def recursively_evaluate(var):
if var is up to date: if var is up to date:
return return
if var.owner.inputs are up to date: if var.owner.inputs are up to date:
update var update var
return return
for input in var.owner.unputs: for input in var.owner.inputs:
recursively_evaluate(var) recursively_evaluate(var)
for output in outputs: for output in outputs:
recursively_evaluate(output) recursively_evaluate(output)
The actual logic is more complex to support intermediate
garbage collection, lazily-evaluated nodes, and better speed. The actual logic is more complex to support intermediate garbage
collection, lazily-evaluated nodes, and better speed.
""" """
...@@ -695,38 +690,36 @@ class Stack(VM): ...@@ -695,38 +690,36 @@ class Stack(VM):
class VMLinker(LocalLinker): class VMLinker(LocalLinker):
""" """Class that satisfies the `Linker` interface by acting as a `VM` factory.
Class that satisfies the Linker interface by acting as a VM factory.
Parameters Parameters
---------- ----------
allow_gc allow_gc
Force the virtual machine to clean up unnecessary Force the virtual machine to clean up unnecessary references, in order
references, in order to allow garbage collection on to allow garbage collection on intermediate values during computation
intermediate values during computation of a function. of a function. If ``None``, use as default the value of the Aesara
If None use as default the value of the Aesara flag allow_gc. flag `allow_gc`.
use_cloop use_cloop
Use the C-based virtual machine if possible Use the C-based virtual machine if possible
callback callback
A callable object to call after each call to a thunk within A callable object to call after each call to a thunk within the virtual
the virtual machine. It will be called with four arguments called machine. It will be called with four arguments: ``node``, ``thunk``,
'node', 'thunk', 'storage_map', and 'compute_map'. ``storage_map``, and ``compute_map``.
callback_input callback_input
A callable object to call on each input to the graph A callable object to call on each input to the graph (variables with no
(variables with no owner). This includes constants and shared owner). This includes constants and shared variables values. It will
variables values. It will be called with two arguments: be called with two arguments: ``var``, ``value``.
'var', 'value'.
lazy lazy
Useful only when use_cloop is False. When lazy is None, use the Useful only when `use_cloop` is False. When `lazy` is ``None``, use the
aesara flag vm__lazy value. Then if we have a None (default) we auto Aesara flag ``vm__lazy`` value. Then if we have a ``None`` (default) we
detect if lazy evaluation is needed and use the appropriate auto detect if lazy evaluation is needed and use the appropriate
version. If lazy is True or False, we force the version used version. If `lazy` is ``True`` or ``False``, we force the version used
between Loop/LoopGC and Stack. between `Loop`/`LoopGC` and `Stack`.
c_thunks c_thunks
If None or True, don't change the default. If False, If ``None`` or ``True``, don't change the default. If ``False``, don't
don't compile c code for the thunks. compile C code for the thunks.
allow_partial_eval allow_partial_eval
If True, enforces usage of Stack or CVM, to allow for partial If ``True``, enforces usage of `Stack` or `CVM`, to allow for partial
evaluation of functions (calculating a subset of outputs). evaluation of functions (calculating a subset of outputs).
""" """
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论