From ca8862a9b2f6b7a96ead3701763333647249091e Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 13 Oct 2022 21:24:16 -0400 Subject: [PATCH 01/17] adding type hints, using keyword args, and start of example based on metropolis alg --- pymc/aesaraf.py | 35 +++++++++++++++++++++++++++------ pymc/smc/smc.py | 4 +++- pymc/step_methods/metropolis.py | 11 +++++++++-- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index ba611d441fe..01f838b5be7 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -57,7 +57,7 @@ ) from aesara.tensor.rewriting.basic import topo_constant_folding from aesara.tensor.rewriting.shape import ShapeFeature -from aesara.tensor.sharedvar import SharedVariable +from aesara.tensor.sharedvar import SharedVariable, TensorSharedVariable from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1 from aesara.tensor.var import TensorConstant, TensorVariable @@ -537,9 +537,9 @@ def join_nonshared_inputs( point: Dict[str, np.ndarray], xs: List[TensorVariable], vars: List[TensorVariable], - shared, + shared: Dict[TensorVariable, TensorSharedVariable], make_shared: bool = False, -): +) -> Tuple[List[TensorVariable], TensorVariable]: """ Takes a list of Aesara Variables and joins their non shared inputs into a single input. @@ -548,11 +548,34 @@ def join_nonshared_inputs( point: a sample point xs: list of Aesara tensors vars: list of variables to join + shared: dict of TensorVariable and their associated TensorSharedVariable. + See return of pm.aesaraf.make_shared_replacements + make_shared: bool + + Examples + -------- + + .. code-block:: python + + initial_point = model.initial_point() + all_model_variables = pm.aesaraf.inputvars(model.value_vars) + # Will be empty here since nothing to replace + shared = pm.aesara.make_shared_replacements( + point=initial_point, + vars=all_model_variables, + model=model, + ) + [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + point=initial_point, + xs=[model.logp()], + vars=all_model_variables, + shared=shared, + ) Returns ------- - tensors, inarray - tensors: list of same tensors but with inarray as input + xs_special, inarray + xs_special: list of same tensors but with inarray as input inarray: vector of inputs """ if not vars: @@ -572,7 +595,7 @@ def join_nonshared_inputs( if aesara.config.compute_test_value != "off": inarray.tag.test_value = joined.tag.test_value - replace = {} + replace: Dict[TensorVariable, TensorSharedVariable] = {} last_idx = 0 for var in vars: shape = point[var.name].shape diff --git a/pymc/smc/smc.py b/pymc/smc/smc.py index 5ff2b9b5cc4..9f12c84d2a4 100644 --- a/pymc/smc/smc.py +++ b/pymc/smc/smc.py @@ -602,7 +602,9 @@ def _logp_forw(point, out_vars, in_vars, shared): out_vars = clone_replace(out_vars, replace_int_input, rebuild_strict=False) in_vars = new_in_vars - out_list, inarray0 = join_nonshared_inputs(point, out_vars, in_vars, shared) + out_list, inarray0 = join_nonshared_inputs( + point=point, xs=out_vars, vars=in_vars, shared=shared + ) f = compile_pymc([inarray0], out_list[0]) f.trust_input = True return f diff --git a/pymc/step_methods/metropolis.py b/pymc/step_methods/metropolis.py index 68909c018c2..91d55108f0e 100644 --- a/pymc/step_methods/metropolis.py +++ b/pymc/step_methods/metropolis.py @@ -13,11 +13,13 @@ # limitations under the License. from typing import Any, Callable, Dict, List, Optional, Tuple +import aesara import numpy as np import numpy.random as nr import scipy.linalg import scipy.special +from aesara import tensor as at from aesara.graph.fg import MissingInputError from aesara.tensor.random.basic import BernoulliRV, CategoricalRV @@ -1052,8 +1054,13 @@ def sample_except(limit, excluded): return candidate -def delta_logp(point, logp, vars, shared): - [logp0], inarray0 = join_nonshared_inputs(point, [logp], vars, shared) +def delta_logp( + point: Dict[str, np.ndarray], + logp: at.TensorVariable, + vars: List[at.TensorVariable], + shared: Dict[at.TensorVariable, at.sharedvar.TensorSharedVariable], +) -> aesara.compile.Function: + [logp0], inarray0 = join_nonshared_inputs(point=point, xs=[logp], vars=vars, shared=shared) tensor_type = inarray0.type inarray1 = tensor_type("inarray1") From d70bb6f3033636b2a65b41b17759f7b36b2292b2 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Wed, 19 Oct 2022 14:48:26 -0400 Subject: [PATCH 02/17] adding second example and some more to the docs --- pymc/aesaraf.py | 44 ++++++++++++++++++++++++++++++-------------- 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 01f838b5be7..7fefe7497a8 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -541,19 +541,26 @@ def join_nonshared_inputs( make_shared: bool = False, ) -> Tuple[List[TensorVariable], TensorVariable]: """ - Takes a list of Aesara Variables and joins their non shared inputs into a single input. + Takes a list of TensorVariables and joins their non shared inputs into a single input. Parameters ---------- - point: a sample point - xs: list of Aesara tensors - vars: list of variables to join - shared: dict of TensorVariable and their associated TensorSharedVariable. - See return of pm.aesaraf.make_shared_replacements - make_shared: bool + point: a sample point associated with a model and variable inputs + xs: list of TensorVariable to replace subgraphs with vars and shared + vars: list of TensorVariable to reshape, join, and return as inarray + shared: dict of TensorVariable and their associated TensorSharedVariable in + subgraph replacement + make_shared: bool flag to use point argument to build subgraphs of xs_special + + Returns + ------- + xs_special, inarray + xs_special: list of same xs TensorVariables but with inarray and shared in subgraphs + inarray: TensorVariable built from vars argument Examples -------- + Join all model variables with none being shared with respect to model logp. .. code-block:: python @@ -572,11 +579,22 @@ def join_nonshared_inputs( shared=shared, ) - Returns - ------- - xs_special, inarray - xs_special: list of same tensors but with inarray as input - inarray: vector of inputs + Same as above but with first model variable being shared. + + .. code-block:: python + + variable_subset = all_model_variable[1:] + shared_subset = pm.aesara.make_shared_replacements( + point=initial_point, + vars=variable_subset, + model=model, + ) + [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + point=initial_point, + xs=[model.logp()], + vars=variable_subset, + shared=shared_subset, + ) """ if not vars: raise ValueError("Empty list of variables.") @@ -587,8 +605,6 @@ def join_nonshared_inputs( tensor_type = joined.type inarray = tensor_type("inarray") else: - if point is None: - raise ValueError("A point is required when `make_shared` is True") joined_values = np.concatenate([point[var.name].ravel() for var in vars]) inarray = aesara.shared(joined_values, "inarray") From 564ab096e40250898e9091dd10cbf3686866a346 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 20 Oct 2022 15:38:03 -0400 Subject: [PATCH 03/17] add the suggestions from PR --- pymc/aesaraf.py | 165 ++++++++++++++++++++++---------- pymc/smc/smc.py | 2 +- pymc/step_methods/metropolis.py | 4 +- 3 files changed, 117 insertions(+), 54 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 7fefe7497a8..c70d6481471 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -535,94 +535,155 @@ def make_shared_replacements(point, vars, model): def join_nonshared_inputs( point: Dict[str, np.ndarray], - xs: List[TensorVariable], - vars: List[TensorVariable], - shared: Dict[TensorVariable, TensorSharedVariable], - make_shared: bool = False, + outputs: List[TensorVariable], + inputs: List[TensorVariable], + shared_inputs: Union[Dict[TensorVariable, TensorSharedVariable], None] = None, + make_inputs_shared: bool = False, ) -> Tuple[List[TensorVariable], TensorVariable]: """ - Takes a list of TensorVariables and joins their non shared inputs into a single input. + Create new outputs and input TensorVariables where the non-shared inputs are joined + in a single raveled vector input. Parameters ---------- - point: a sample point associated with a model and variable inputs - xs: list of TensorVariable to replace subgraphs with vars and shared - vars: list of TensorVariable to reshape, join, and return as inarray - shared: dict of TensorVariable and their associated TensorSharedVariable in - subgraph replacement - make_shared: bool flag to use point argument to build subgraphs of xs_special + point: dictionary of variable names to numerical values + Dictionary that maps each input variable name to a numerical variable. The values + are used to extract the shape of each input variable to establish a correct + mapping between joined and original inputs. The shape of each variable is + assumed to be fixed. + outputs: list of TensorVariable + List of output TensorVariables whose non-shared inputs will be replaced + by a joined vector input. + inputs: list of TensorVariable + List of input TensorVariables which will be replaced by a joined vector input. + shared_inputs: dict of TensorVariable to TensorSharedVariable, optional + Dict of TensorVariable and their associated TensorSharedVariable in + subgraph replacement. + make_inputs_shared: bool, defaults to False + Whether to make the joined vector input a shared variable. Returns ------- - xs_special, inarray - xs_special: list of same xs TensorVariables but with inarray and shared in subgraphs - inarray: TensorVariable built from vars argument + new_outputs, joined_inputs + new_outputs: list of TensorVariable + List of new outputs `outputs` TensorVariables that depend on `joined_inputs` and new shared variables as inputs. + joined_inputs: TensorVariable + Joined input vector TensorVariable for the `new_outputs` Examples -------- - Join all model variables with none being shared with respect to model logp. + Join the inputs of a simple Aesara graph .. code-block:: python + import aesara.tensor as at + import numpy as np + + from pymc.aesaraf import join_nonshared_inputs + + # Original non-shared inputs + x = at.scalar("x") + y = at.vector("y") + # Orignal output + out = x + y + print(out.eval({x: np.array(1), y: np.array([1, 2, 3])})) # [2, 3, 4] + + # New output and inputs + [new_out], joined_inputs = join_nonshared_inputs( + point={ + "x": np.zeros(()), + "y": np.zeros(3), + }, + outputs=[out], + inputs=[x, y], + ) + print(new_out.eval({ + joined_inputs: np.array([1, 1, 2, 3]), + })) # [2, 3, 4] + + Join the input value variables of a model logp. + + .. code-block:: python + + import pymc as pm + + with pm.Model() as model: + mu_pop = pm.Normal("mu_pop") + sigma_pop = pm.HalfNormal("sigma_pop") + mu = pm.Normal("mu", mu_pop, sigma_pop, shape=(3, )) + + y = pm.Normal("y", mu, 1.0, observed=[0, 1, 2]) + + print(model.compile_logp()({ + "mu_pop": 0, + "sigma_pop_log__": 1, + "mu": [0, 1, 2], + })) # -12.691227342634292 initial_point = model.initial_point() - all_model_variables = pm.aesaraf.inputvars(model.value_vars) - # Will be empty here since nothing to replace - shared = pm.aesara.make_shared_replacements( - point=initial_point, - vars=all_model_variables, - model=model, - ) - [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + inputs = model.value_vars + + [logp], joined_inputs = join_nonshared_inputs( point=initial_point, - xs=[model.logp()], - vars=all_model_variables, - shared=shared, + outputs=[model.logp()], + inputs=inputs, ) - Same as above but with first model variable being shared. + print(logp.eval({ + joined_inputs: [0, 1, 0, 1, 2], + })) + + Same as above but with the `mu_pop` value variable being shared. .. code-block:: python - variable_subset = all_model_variable[1:] - shared_subset = pm.aesara.make_shared_replacements( - point=initial_point, - vars=variable_subset, - model=model, - ) - [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + from aesara import shared + + mu_pop_input, *other_inputs = inputs + shared_mu_pop_input = shared(0.0) + + [logp], other_joined_inputs = join_nonshared_inputs( point=initial_point, - xs=[model.logp()], - vars=variable_subset, - shared=shared_subset, + outputs=[model.logp()], + inputs=other_inputs, + shared_inputs={ + mu_pop_input: shared_mu_pop_input + }, ) + + print(logp.eval({ + other_joined_inputs: [1, 0, 1, 2], + })) # -12.691227342634292 """ - if not vars: - raise ValueError("Empty list of variables.") + if not inputs: + raise ValueError("Empty list of input variables.") - joined = at.concatenate([var.ravel() for var in vars]) + raveled_inputs = at.concatenate([var.ravel() for var in inputs]) - if not make_shared: - tensor_type = joined.type - inarray = tensor_type("inarray") + if not make_inputs_shared: + tensor_type = raveled_inputs.type + joined_inputs = tensor_type("joined_inputs") else: - joined_values = np.concatenate([point[var.name].ravel() for var in vars]) - inarray = aesara.shared(joined_values, "inarray") + joined_values = np.concatenate([point[var.name].ravel() for var in inputs]) + joined_inputs = aesara.shared(joined_values, "joined_inputs") if aesara.config.compute_test_value != "off": - inarray.tag.test_value = joined.tag.test_value + joined_inputs.tag.test_value = raveled_inputs.tag.test_value - replace: Dict[TensorVariable, TensorSharedVariable] = {} + replace: Dict[TensorVariable, TensorVariable] = {} last_idx = 0 - for var in vars: + for var in inputs: shape = point[var.name].shape arr_len = np.prod(shape, dtype=int) - replace[var] = inarray[last_idx : last_idx + arr_len].reshape(shape).astype(var.dtype) + replace[var] = joined_inputs[last_idx : last_idx + arr_len].reshape(shape).astype(var.dtype) last_idx += arr_len - replace.update(shared) + if shared_inputs is not None: + replace.update(shared_inputs) - xs_special = [aesara.clone_replace(x, replace, rebuild_strict=False) for x in xs] - return xs_special, inarray + new_outputs = [ + aesara.clone_replace(output, replace, rebuild_strict=False) for output in outputs + ] + return new_outputs, joined_inputs class PointFunc: diff --git a/pymc/smc/smc.py b/pymc/smc/smc.py index 9f12c84d2a4..58b06aa8211 100644 --- a/pymc/smc/smc.py +++ b/pymc/smc/smc.py @@ -603,7 +603,7 @@ def _logp_forw(point, out_vars, in_vars, shared): in_vars = new_in_vars out_list, inarray0 = join_nonshared_inputs( - point=point, xs=out_vars, vars=in_vars, shared=shared + point=point, outputs=out_vars, inputs=in_vars, shared_inputs=shared ) f = compile_pymc([inarray0], out_list[0]) f.trust_input = True diff --git a/pymc/step_methods/metropolis.py b/pymc/step_methods/metropolis.py index 91d55108f0e..91b475652c9 100644 --- a/pymc/step_methods/metropolis.py +++ b/pymc/step_methods/metropolis.py @@ -1060,7 +1060,9 @@ def delta_logp( vars: List[at.TensorVariable], shared: Dict[at.TensorVariable, at.sharedvar.TensorSharedVariable], ) -> aesara.compile.Function: - [logp0], inarray0 = join_nonshared_inputs(point=point, xs=[logp], vars=vars, shared=shared) + [logp0], inarray0 = join_nonshared_inputs( + point=point, outputs=[logp], inputs=vars, shared_inputs=shared + ) tensor_type = inarray0.type inarray1 = tensor_type("inarray1") From a526a03b32eb4b52dfc924f9c74478224adf4a6e Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 20 Oct 2022 16:05:19 -0400 Subject: [PATCH 04/17] first example did not render --- pymc/aesaraf.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index c70d6481471..923659ade18 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -572,9 +572,10 @@ def join_nonshared_inputs( Examples -------- - Join the inputs of a simple Aesara graph + Join the inputs of a simple Aesara graph. .. code-block:: python + import aesara.tensor as at import numpy as np From f9a165aa05a6cd6c4ee595c377297bfe8fb47d70 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Fri, 21 Oct 2022 08:05:42 -0400 Subject: [PATCH 05/17] fix typo --- pymc/aesaraf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 923659ade18..67d27eb9bb6 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -584,7 +584,7 @@ def join_nonshared_inputs( # Original non-shared inputs x = at.scalar("x") y = at.vector("y") - # Orignal output + # Original output out = x + y print(out.eval({x: np.array(1), y: np.array([1, 2, 3])})) # [2, 3, 4] From c4be3230197f185a8905b54298c80f38cf0ee23c Mon Sep 17 00:00:00 2001 From: Will Dean Date: Wed, 26 Oct 2022 18:50:06 -0400 Subject: [PATCH 06/17] switch to numpy docstyle. comment in example. fix return --- pymc/aesaraf.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 67d27eb9bb6..2d4a9ac3a7c 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -537,7 +537,7 @@ def join_nonshared_inputs( point: Dict[str, np.ndarray], outputs: List[TensorVariable], inputs: List[TensorVariable], - shared_inputs: Union[Dict[TensorVariable, TensorSharedVariable], None] = None, + shared_inputs: Optional[Dict[TensorVariable, TensorSharedVariable]] = None, make_inputs_shared: bool = False, ) -> Tuple[List[TensorVariable], TensorVariable]: """ @@ -546,7 +546,7 @@ def join_nonshared_inputs( Parameters ---------- - point: dictionary of variable names to numerical values + point: dict of {str : array_like} Dictionary that maps each input variable name to a numerical variable. The values are used to extract the shape of each input variable to establish a correct mapping between joined and original inputs. The shape of each variable is @@ -564,7 +564,6 @@ def join_nonshared_inputs( Returns ------- - new_outputs, joined_inputs new_outputs: list of TensorVariable List of new outputs `outputs` TensorVariables that depend on `joined_inputs` and new shared variables as inputs. joined_inputs: TensorVariable @@ -590,7 +589,7 @@ def join_nonshared_inputs( # New output and inputs [new_out], joined_inputs = join_nonshared_inputs( - point={ + point={ # Only shapes matter "x": np.zeros(()), "y": np.zeros(3), }, @@ -631,7 +630,7 @@ def join_nonshared_inputs( print(logp.eval({ joined_inputs: [0, 1, 0, 1, 2], - })) + })) # -12.691227342634292 Same as above but with the `mu_pop` value variable being shared. From b3ca27d8088636d17be50ae23b993bedcb0fde61 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Wed, 26 Oct 2022 22:17:33 -0400 Subject: [PATCH 07/17] forgotten dict --- pymc/aesaraf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 17815d1eeee..7309a8d4452 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -556,7 +556,7 @@ def join_nonshared_inputs( by a joined vector input. inputs: list of TensorVariable List of input TensorVariables which will be replaced by a joined vector input. - shared_inputs: dict of TensorVariable to TensorSharedVariable, optional + shared_inputs: dict of {TensorVariable : TensorSharedVariable}, optional Dict of TensorVariable and their associated TensorSharedVariable in subgraph replacement. make_inputs_shared: bool, defaults to False From 7e193503a93c66545316ae30b86f4d948c66a3e2 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 27 Oct 2022 08:50:48 -0400 Subject: [PATCH 08/17] numpy docstring standards --- pymc/aesaraf.py | 14 +++++++------- pymc/smc/smc.py | 6 +++--- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 7309a8d4452..020216278f9 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -546,27 +546,27 @@ def join_nonshared_inputs( Parameters ---------- - point: dict of {str : array_like} + point : Point Dictionary that maps each input variable name to a numerical variable. The values are used to extract the shape of each input variable to establish a correct mapping between joined and original inputs. The shape of each variable is assumed to be fixed. - outputs: list of TensorVariable + outputs : list of TensorVariable List of output TensorVariables whose non-shared inputs will be replaced by a joined vector input. - inputs: list of TensorVariable + inputs : list of TensorVariable List of input TensorVariables which will be replaced by a joined vector input. - shared_inputs: dict of {TensorVariable : TensorSharedVariable}, optional + shared_inputs : dict of {TensorVariable : TensorSharedVariable}, optional Dict of TensorVariable and their associated TensorSharedVariable in subgraph replacement. - make_inputs_shared: bool, defaults to False + make_inputs_shared : bool, default False Whether to make the joined vector input a shared variable. Returns ------- - new_outputs: list of TensorVariable + new_outputs : list of TensorVariable List of new outputs `outputs` TensorVariables that depend on `joined_inputs` and new shared variables as inputs. - joined_inputs: TensorVariable + joined_inputs : TensorVariable Joined input vector TensorVariable for the `new_outputs` Examples diff --git a/pymc/smc/smc.py b/pymc/smc/smc.py index 58b06aa8211..5cd6592b368 100644 --- a/pymc/smc/smc.py +++ b/pymc/smc/smc.py @@ -579,11 +579,11 @@ def _logp_forw(point, out_vars, in_vars, shared): Parameters ---------- - out_vars: List + out_vars : list containing :class:`pymc.Distribution` for the output variables - in_vars: List + in_vars : list containing :class:`pymc.Distribution` for the input variables - shared: List + shared : list containing :class:`aesara.tensor.Tensor` for depended shared data """ From 2923927c79a3c77089f6906f86a5a5ff09d8b834 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 13 Oct 2022 21:24:16 -0400 Subject: [PATCH 09/17] adding type hints, using keyword args, and start of example based on metropolis alg --- pymc/aesaraf.py | 35 +++++++++++++++++++++++++++------ pymc/smc/kernels.py | 4 +++- pymc/step_methods/metropolis.py | 11 +++++++++-- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 84be5d37197..4f8bd8db8fc 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -57,7 +57,7 @@ ) from aesara.tensor.rewriting.basic import topo_constant_folding from aesara.tensor.rewriting.shape import ShapeFeature -from aesara.tensor.sharedvar import SharedVariable +from aesara.tensor.sharedvar import SharedVariable, TensorSharedVariable from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1 from aesara.tensor.var import TensorConstant, TensorVariable @@ -537,9 +537,9 @@ def join_nonshared_inputs( point: Dict[str, np.ndarray], xs: List[TensorVariable], vars: List[TensorVariable], - shared, + shared: Dict[TensorVariable, TensorSharedVariable], make_shared: bool = False, -): +) -> Tuple[List[TensorVariable], TensorVariable]: """ Takes a list of Aesara Variables and joins their non shared inputs into a single input. @@ -548,11 +548,34 @@ def join_nonshared_inputs( point: a sample point xs: list of Aesara tensors vars: list of variables to join + shared: dict of TensorVariable and their associated TensorSharedVariable. + See return of pm.aesaraf.make_shared_replacements + make_shared: bool + + Examples + -------- + + .. code-block:: python + + initial_point = model.initial_point() + all_model_variables = pm.aesaraf.inputvars(model.value_vars) + # Will be empty here since nothing to replace + shared = pm.aesara.make_shared_replacements( + point=initial_point, + vars=all_model_variables, + model=model, + ) + [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + point=initial_point, + xs=[model.logp()], + vars=all_model_variables, + shared=shared, + ) Returns ------- - tensors, inarray - tensors: list of same tensors but with inarray as input + xs_special, inarray + xs_special: list of same tensors but with inarray as input inarray: vector of inputs """ if not vars: @@ -572,7 +595,7 @@ def join_nonshared_inputs( if aesara.config.compute_test_value != "off": inarray.tag.test_value = joined.tag.test_value - replace = {} + replace: Dict[TensorVariable, TensorSharedVariable] = {} last_idx = 0 for var in vars: shape = point[var.name].shape diff --git a/pymc/smc/kernels.py b/pymc/smc/kernels.py index 5ff2b9b5cc4..9f12c84d2a4 100644 --- a/pymc/smc/kernels.py +++ b/pymc/smc/kernels.py @@ -602,7 +602,9 @@ def _logp_forw(point, out_vars, in_vars, shared): out_vars = clone_replace(out_vars, replace_int_input, rebuild_strict=False) in_vars = new_in_vars - out_list, inarray0 = join_nonshared_inputs(point, out_vars, in_vars, shared) + out_list, inarray0 = join_nonshared_inputs( + point=point, xs=out_vars, vars=in_vars, shared=shared + ) f = compile_pymc([inarray0], out_list[0]) f.trust_input = True return f diff --git a/pymc/step_methods/metropolis.py b/pymc/step_methods/metropolis.py index 42f2a6cdb12..c6876317e12 100644 --- a/pymc/step_methods/metropolis.py +++ b/pymc/step_methods/metropolis.py @@ -13,11 +13,13 @@ # limitations under the License. from typing import Any, Callable, Dict, List, Optional, Tuple +import aesara import numpy as np import numpy.random as nr import scipy.linalg import scipy.special +from aesara import tensor as at from aesara.graph.fg import MissingInputError from aesara.tensor.random.basic import BernoulliRV, CategoricalRV @@ -1052,8 +1054,13 @@ def sample_except(limit, excluded): return candidate -def delta_logp(point, logp, vars, shared): - [logp0], inarray0 = join_nonshared_inputs(point, [logp], vars, shared) +def delta_logp( + point: Dict[str, np.ndarray], + logp: at.TensorVariable, + vars: List[at.TensorVariable], + shared: Dict[at.TensorVariable, at.sharedvar.TensorSharedVariable], +) -> aesara.compile.Function: + [logp0], inarray0 = join_nonshared_inputs(point=point, xs=[logp], vars=vars, shared=shared) tensor_type = inarray0.type inarray1 = tensor_type("inarray1") From 3a8b3fb996d0af8a8328c399a0d1adceb77b7927 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Wed, 19 Oct 2022 14:48:26 -0400 Subject: [PATCH 10/17] adding second example and some more to the docs --- pymc/aesaraf.py | 44 ++++++++++++++++++++++++++++++-------------- 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 4f8bd8db8fc..e42137818b5 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -541,19 +541,26 @@ def join_nonshared_inputs( make_shared: bool = False, ) -> Tuple[List[TensorVariable], TensorVariable]: """ - Takes a list of Aesara Variables and joins their non shared inputs into a single input. + Takes a list of TensorVariables and joins their non shared inputs into a single input. Parameters ---------- - point: a sample point - xs: list of Aesara tensors - vars: list of variables to join - shared: dict of TensorVariable and their associated TensorSharedVariable. - See return of pm.aesaraf.make_shared_replacements - make_shared: bool + point: a sample point associated with a model and variable inputs + xs: list of TensorVariable to replace subgraphs with vars and shared + vars: list of TensorVariable to reshape, join, and return as inarray + shared: dict of TensorVariable and their associated TensorSharedVariable in + subgraph replacement + make_shared: bool flag to use point argument to build subgraphs of xs_special + + Returns + ------- + xs_special, inarray + xs_special: list of same xs TensorVariables but with inarray and shared in subgraphs + inarray: TensorVariable built from vars argument Examples -------- + Join all model variables with none being shared with respect to model logp. .. code-block:: python @@ -572,11 +579,22 @@ def join_nonshared_inputs( shared=shared, ) - Returns - ------- - xs_special, inarray - xs_special: list of same tensors but with inarray as input - inarray: vector of inputs + Same as above but with first model variable being shared. + + .. code-block:: python + + variable_subset = all_model_variable[1:] + shared_subset = pm.aesara.make_shared_replacements( + point=initial_point, + vars=variable_subset, + model=model, + ) + [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + point=initial_point, + xs=[model.logp()], + vars=variable_subset, + shared=shared_subset, + ) """ if not vars: raise ValueError("Empty list of variables.") @@ -587,8 +605,6 @@ def join_nonshared_inputs( tensor_type = joined.type inarray = tensor_type("inarray") else: - if point is None: - raise ValueError("A point is required when `make_shared` is True") joined_values = np.concatenate([point[var.name].ravel() for var in vars]) inarray = aesara.shared(joined_values, "inarray") From 70601f549f16898c2a61853d93faf75055b16bcc Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 20 Oct 2022 15:38:03 -0400 Subject: [PATCH 11/17] add the suggestions from PR --- pymc/aesaraf.py | 165 ++++++++++++++++++++++---------- pymc/smc/kernels.py | 2 +- pymc/step_methods/metropolis.py | 4 +- 3 files changed, 117 insertions(+), 54 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index e42137818b5..d1d2d38a2b1 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -535,94 +535,155 @@ def make_shared_replacements(point, vars, model): def join_nonshared_inputs( point: Dict[str, np.ndarray], - xs: List[TensorVariable], - vars: List[TensorVariable], - shared: Dict[TensorVariable, TensorSharedVariable], - make_shared: bool = False, + outputs: List[TensorVariable], + inputs: List[TensorVariable], + shared_inputs: Union[Dict[TensorVariable, TensorSharedVariable], None] = None, + make_inputs_shared: bool = False, ) -> Tuple[List[TensorVariable], TensorVariable]: """ - Takes a list of TensorVariables and joins their non shared inputs into a single input. + Create new outputs and input TensorVariables where the non-shared inputs are joined + in a single raveled vector input. Parameters ---------- - point: a sample point associated with a model and variable inputs - xs: list of TensorVariable to replace subgraphs with vars and shared - vars: list of TensorVariable to reshape, join, and return as inarray - shared: dict of TensorVariable and their associated TensorSharedVariable in - subgraph replacement - make_shared: bool flag to use point argument to build subgraphs of xs_special + point: dictionary of variable names to numerical values + Dictionary that maps each input variable name to a numerical variable. The values + are used to extract the shape of each input variable to establish a correct + mapping between joined and original inputs. The shape of each variable is + assumed to be fixed. + outputs: list of TensorVariable + List of output TensorVariables whose non-shared inputs will be replaced + by a joined vector input. + inputs: list of TensorVariable + List of input TensorVariables which will be replaced by a joined vector input. + shared_inputs: dict of TensorVariable to TensorSharedVariable, optional + Dict of TensorVariable and their associated TensorSharedVariable in + subgraph replacement. + make_inputs_shared: bool, defaults to False + Whether to make the joined vector input a shared variable. Returns ------- - xs_special, inarray - xs_special: list of same xs TensorVariables but with inarray and shared in subgraphs - inarray: TensorVariable built from vars argument + new_outputs, joined_inputs + new_outputs: list of TensorVariable + List of new outputs `outputs` TensorVariables that depend on `joined_inputs` and new shared variables as inputs. + joined_inputs: TensorVariable + Joined input vector TensorVariable for the `new_outputs` Examples -------- - Join all model variables with none being shared with respect to model logp. + Join the inputs of a simple Aesara graph .. code-block:: python + import aesara.tensor as at + import numpy as np + + from pymc.aesaraf import join_nonshared_inputs + + # Original non-shared inputs + x = at.scalar("x") + y = at.vector("y") + # Orignal output + out = x + y + print(out.eval({x: np.array(1), y: np.array([1, 2, 3])})) # [2, 3, 4] + + # New output and inputs + [new_out], joined_inputs = join_nonshared_inputs( + point={ + "x": np.zeros(()), + "y": np.zeros(3), + }, + outputs=[out], + inputs=[x, y], + ) + print(new_out.eval({ + joined_inputs: np.array([1, 1, 2, 3]), + })) # [2, 3, 4] + + Join the input value variables of a model logp. + + .. code-block:: python + + import pymc as pm + + with pm.Model() as model: + mu_pop = pm.Normal("mu_pop") + sigma_pop = pm.HalfNormal("sigma_pop") + mu = pm.Normal("mu", mu_pop, sigma_pop, shape=(3, )) + + y = pm.Normal("y", mu, 1.0, observed=[0, 1, 2]) + + print(model.compile_logp()({ + "mu_pop": 0, + "sigma_pop_log__": 1, + "mu": [0, 1, 2], + })) # -12.691227342634292 initial_point = model.initial_point() - all_model_variables = pm.aesaraf.inputvars(model.value_vars) - # Will be empty here since nothing to replace - shared = pm.aesara.make_shared_replacements( - point=initial_point, - vars=all_model_variables, - model=model, - ) - [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + inputs = model.value_vars + + [logp], joined_inputs = join_nonshared_inputs( point=initial_point, - xs=[model.logp()], - vars=all_model_variables, - shared=shared, + outputs=[model.logp()], + inputs=inputs, ) - Same as above but with first model variable being shared. + print(logp.eval({ + joined_inputs: [0, 1, 0, 1, 2], + })) + + Same as above but with the `mu_pop` value variable being shared. .. code-block:: python - variable_subset = all_model_variable[1:] - shared_subset = pm.aesara.make_shared_replacements( - point=initial_point, - vars=variable_subset, - model=model, - ) - [logp0], inarray0 = pm.aesara.join_nonshared_inputs( + from aesara import shared + + mu_pop_input, *other_inputs = inputs + shared_mu_pop_input = shared(0.0) + + [logp], other_joined_inputs = join_nonshared_inputs( point=initial_point, - xs=[model.logp()], - vars=variable_subset, - shared=shared_subset, + outputs=[model.logp()], + inputs=other_inputs, + shared_inputs={ + mu_pop_input: shared_mu_pop_input + }, ) + + print(logp.eval({ + other_joined_inputs: [1, 0, 1, 2], + })) # -12.691227342634292 """ - if not vars: - raise ValueError("Empty list of variables.") + if not inputs: + raise ValueError("Empty list of input variables.") - joined = at.concatenate([var.ravel() for var in vars]) + raveled_inputs = at.concatenate([var.ravel() for var in inputs]) - if not make_shared: - tensor_type = joined.type - inarray = tensor_type("inarray") + if not make_inputs_shared: + tensor_type = raveled_inputs.type + joined_inputs = tensor_type("joined_inputs") else: - joined_values = np.concatenate([point[var.name].ravel() for var in vars]) - inarray = aesara.shared(joined_values, "inarray") + joined_values = np.concatenate([point[var.name].ravel() for var in inputs]) + joined_inputs = aesara.shared(joined_values, "joined_inputs") if aesara.config.compute_test_value != "off": - inarray.tag.test_value = joined.tag.test_value + joined_inputs.tag.test_value = raveled_inputs.tag.test_value - replace: Dict[TensorVariable, TensorSharedVariable] = {} + replace: Dict[TensorVariable, TensorVariable] = {} last_idx = 0 - for var in vars: + for var in inputs: shape = point[var.name].shape arr_len = np.prod(shape, dtype=int) - replace[var] = inarray[last_idx : last_idx + arr_len].reshape(shape).astype(var.dtype) + replace[var] = joined_inputs[last_idx : last_idx + arr_len].reshape(shape).astype(var.dtype) last_idx += arr_len - replace.update(shared) + if shared_inputs is not None: + replace.update(shared_inputs) - xs_special = [aesara.clone_replace(x, replace, rebuild_strict=False) for x in xs] - return xs_special, inarray + new_outputs = [ + aesara.clone_replace(output, replace, rebuild_strict=False) for output in outputs + ] + return new_outputs, joined_inputs class PointFunc: diff --git a/pymc/smc/kernels.py b/pymc/smc/kernels.py index 9f12c84d2a4..58b06aa8211 100644 --- a/pymc/smc/kernels.py +++ b/pymc/smc/kernels.py @@ -603,7 +603,7 @@ def _logp_forw(point, out_vars, in_vars, shared): in_vars = new_in_vars out_list, inarray0 = join_nonshared_inputs( - point=point, xs=out_vars, vars=in_vars, shared=shared + point=point, outputs=out_vars, inputs=in_vars, shared_inputs=shared ) f = compile_pymc([inarray0], out_list[0]) f.trust_input = True diff --git a/pymc/step_methods/metropolis.py b/pymc/step_methods/metropolis.py index c6876317e12..00b9ecd7525 100644 --- a/pymc/step_methods/metropolis.py +++ b/pymc/step_methods/metropolis.py @@ -1060,7 +1060,9 @@ def delta_logp( vars: List[at.TensorVariable], shared: Dict[at.TensorVariable, at.sharedvar.TensorSharedVariable], ) -> aesara.compile.Function: - [logp0], inarray0 = join_nonshared_inputs(point=point, xs=[logp], vars=vars, shared=shared) + [logp0], inarray0 = join_nonshared_inputs( + point=point, outputs=[logp], inputs=vars, shared_inputs=shared + ) tensor_type = inarray0.type inarray1 = tensor_type("inarray1") From acc4f6db6758325664fc5ad4d7d8056047e32df0 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 20 Oct 2022 16:05:19 -0400 Subject: [PATCH 12/17] first example did not render --- pymc/aesaraf.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index d1d2d38a2b1..aa33e497ce2 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -572,9 +572,10 @@ def join_nonshared_inputs( Examples -------- - Join the inputs of a simple Aesara graph + Join the inputs of a simple Aesara graph. .. code-block:: python + import aesara.tensor as at import numpy as np From fb0a1419c4bd8abd9734a0d1f2ae785311f5b494 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Fri, 21 Oct 2022 08:05:42 -0400 Subject: [PATCH 13/17] fix typo --- pymc/aesaraf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index aa33e497ce2..4e536344ec3 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -584,7 +584,7 @@ def join_nonshared_inputs( # Original non-shared inputs x = at.scalar("x") y = at.vector("y") - # Orignal output + # Original output out = x + y print(out.eval({x: np.array(1), y: np.array([1, 2, 3])})) # [2, 3, 4] From 4f9d68781043b330684dda6d132c9bfa5215bd3f Mon Sep 17 00:00:00 2001 From: Will Dean Date: Wed, 26 Oct 2022 18:50:06 -0400 Subject: [PATCH 14/17] switch to numpy docstyle. comment in example. fix return --- pymc/aesaraf.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 4e536344ec3..17815d1eeee 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -537,7 +537,7 @@ def join_nonshared_inputs( point: Dict[str, np.ndarray], outputs: List[TensorVariable], inputs: List[TensorVariable], - shared_inputs: Union[Dict[TensorVariable, TensorSharedVariable], None] = None, + shared_inputs: Optional[Dict[TensorVariable, TensorSharedVariable]] = None, make_inputs_shared: bool = False, ) -> Tuple[List[TensorVariable], TensorVariable]: """ @@ -546,7 +546,7 @@ def join_nonshared_inputs( Parameters ---------- - point: dictionary of variable names to numerical values + point: dict of {str : array_like} Dictionary that maps each input variable name to a numerical variable. The values are used to extract the shape of each input variable to establish a correct mapping between joined and original inputs. The shape of each variable is @@ -564,7 +564,6 @@ def join_nonshared_inputs( Returns ------- - new_outputs, joined_inputs new_outputs: list of TensorVariable List of new outputs `outputs` TensorVariables that depend on `joined_inputs` and new shared variables as inputs. joined_inputs: TensorVariable @@ -590,7 +589,7 @@ def join_nonshared_inputs( # New output and inputs [new_out], joined_inputs = join_nonshared_inputs( - point={ + point={ # Only shapes matter "x": np.zeros(()), "y": np.zeros(3), }, @@ -631,7 +630,7 @@ def join_nonshared_inputs( print(logp.eval({ joined_inputs: [0, 1, 0, 1, 2], - })) + })) # -12.691227342634292 Same as above but with the `mu_pop` value variable being shared. From a1dc3c493daf0aade6364f274fcbce9fb73a1fe2 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Wed, 26 Oct 2022 22:17:33 -0400 Subject: [PATCH 15/17] forgotten dict --- pymc/aesaraf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 17815d1eeee..7309a8d4452 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -556,7 +556,7 @@ def join_nonshared_inputs( by a joined vector input. inputs: list of TensorVariable List of input TensorVariables which will be replaced by a joined vector input. - shared_inputs: dict of TensorVariable to TensorSharedVariable, optional + shared_inputs: dict of {TensorVariable : TensorSharedVariable}, optional Dict of TensorVariable and their associated TensorSharedVariable in subgraph replacement. make_inputs_shared: bool, defaults to False From 70659820ef6a94001bf3d562bd31f1193b40b139 Mon Sep 17 00:00:00 2001 From: Will Dean Date: Thu, 27 Oct 2022 08:50:48 -0400 Subject: [PATCH 16/17] numpy docstring standards --- pymc/aesaraf.py | 14 +++++++------- pymc/smc/kernels.py | 6 +++--- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 7309a8d4452..020216278f9 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -546,27 +546,27 @@ def join_nonshared_inputs( Parameters ---------- - point: dict of {str : array_like} + point : Point Dictionary that maps each input variable name to a numerical variable. The values are used to extract the shape of each input variable to establish a correct mapping between joined and original inputs. The shape of each variable is assumed to be fixed. - outputs: list of TensorVariable + outputs : list of TensorVariable List of output TensorVariables whose non-shared inputs will be replaced by a joined vector input. - inputs: list of TensorVariable + inputs : list of TensorVariable List of input TensorVariables which will be replaced by a joined vector input. - shared_inputs: dict of {TensorVariable : TensorSharedVariable}, optional + shared_inputs : dict of {TensorVariable : TensorSharedVariable}, optional Dict of TensorVariable and their associated TensorSharedVariable in subgraph replacement. - make_inputs_shared: bool, defaults to False + make_inputs_shared : bool, default False Whether to make the joined vector input a shared variable. Returns ------- - new_outputs: list of TensorVariable + new_outputs : list of TensorVariable List of new outputs `outputs` TensorVariables that depend on `joined_inputs` and new shared variables as inputs. - joined_inputs: TensorVariable + joined_inputs : TensorVariable Joined input vector TensorVariable for the `new_outputs` Examples diff --git a/pymc/smc/kernels.py b/pymc/smc/kernels.py index 58b06aa8211..5cd6592b368 100644 --- a/pymc/smc/kernels.py +++ b/pymc/smc/kernels.py @@ -579,11 +579,11 @@ def _logp_forw(point, out_vars, in_vars, shared): Parameters ---------- - out_vars: List + out_vars : list containing :class:`pymc.Distribution` for the output variables - in_vars: List + in_vars : list containing :class:`pymc.Distribution` for the input variables - shared: List + shared : list containing :class:`aesara.tensor.Tensor` for depended shared data """ From a50be1f9b1460e7773117a620834de743cf2fb4b Mon Sep 17 00:00:00 2001 From: Will Dean Date: Tue, 1 Nov 2022 18:09:56 -0400 Subject: [PATCH 17/17] switching away from Point --- pymc/aesaraf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/aesaraf.py b/pymc/aesaraf.py index 020216278f9..a22fa5abc9d 100644 --- a/pymc/aesaraf.py +++ b/pymc/aesaraf.py @@ -546,7 +546,7 @@ def join_nonshared_inputs( Parameters ---------- - point : Point + point : dict of {str : array_like} Dictionary that maps each input variable name to a numerical variable. The values are used to extract the shape of each input variable to establish a correct mapping between joined and original inputs. The shape of each variable is