From d4cdf3304619cbe7c3891b7b30b548b3643515ff Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Sun, 16 Jun 2024 18:13:02 -0500 Subject: [PATCH 1/9] Started adding usage examples to pytensor math functions --- pytensor/tensor/basic.py | 26 +++++++++++++++++++++++++- pytensor/tensor/math.py | 12 ++++++++++++ pytensor/tensor/special.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index e8964ccdf9..d31b277cd8 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -757,8 +757,32 @@ def cast(x, dtype: str | np.dtype) -> TensorVariable: @scalar_elemwise def switch(cond, ift, iff): - """if cond then ift else iff""" + """ + Conditionally selects elements from two tensors based on a condition tensor. + + This op is similar to NumPy's `np.where` and `np.choose` functions. + + Parameters + ---------- + cond : TensorVariable + A boolean-type tensor determining which output value to choose. + Should be broadcastable to the shapes of `ift` and `iff`. + ift : TensorVariable + Values selected at `True` elements of `cond`. + iff : TensorVariable + Values selected at `False` elements of `cond`. + Examples + -------- + This example demonstrates how `switch` can be used in PyMC to model a + categorical variable. + + .. code:: python + + with pm.Model(): + x = pm.Categorical('x', np.array([0.1, 0.9])) + y_ = pm.Bernoulli('y_', p=switch(x, 0.9, 0.1), shape=10) + """ where = switch diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index d515d51c3a..5cfac73c1e 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -1945,6 +1945,18 @@ def dot(l, r): """Return a symbolic dot product. This is designed to work with both sparse and dense tensors types. + + Example usage with PyMC: + + .. code:: python + + import pymc3 as pm + import theano.tensor as tt + + with pm.Model() as model: + x = pm.Normal('x', mu=0, sd=1, shape=2) + y = pm.Normal('y', mu=0, sd=1, shape=2) + z = tt.dot(x, y) """ if not isinstance(l, Variable): diff --git a/pytensor/tensor/special.py b/pytensor/tensor/special.py index a2f02fabd8..b743c1dadb 100644 --- a/pytensor/tensor/special.py +++ b/pytensor/tensor/special.py @@ -485,6 +485,36 @@ def c_code_cache_version(): def softmax(c, axis=None): + """ + Compute the softmax of a vector along a specified axis. + + Parameters + ---------- + c : TensorVariable + The input tensor. + axis : int + The axis along which to compute the softmax. + + Returns + ------- + TensorVariable + The softmax of the input tensor along the specified axis. + + Examples + -------- + In PyMC, you can use this function to compute a softmax over a vector of + probabilities representing the likelihood of each class in a multiclass + classification problem. Here is an example:: + + import pymc as pm + import pytensor.tensor as pt + + with pm.Model() as model: + weights = pm.Gamma('weights', 1, 1, shape=3) + softmax_prob = pt.softmax(weights) + outcome = pm.Categorical('outcome', p=softmax_prob) + + """ c = as_tensor_variable(c) return Softmax(axis=axis)(c) From 7fe1c81670c0ed7c6849b2f3a903a335bf358006 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 08:13:14 -0500 Subject: [PATCH 2/9] Fixed docstring typo --- pytensor/tensor/math.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 5cfac73c1e..a5668d14d3 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -1950,13 +1950,13 @@ def dot(l, r): .. code:: python - import pymc3 as pm - import theano.tensor as tt + import pymc as pm + import pytensor.tensor as pt with pm.Model() as model: x = pm.Normal('x', mu=0, sd=1, shape=2) y = pm.Normal('y', mu=0, sd=1, shape=2) - z = tt.dot(x, y) + z = pt.dot(x, y) """ if not isinstance(l, Variable): From bba6237a621a5d110a0675b0c27ed368f7fb9eaa Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 08:18:14 -0500 Subject: [PATCH 3/9] Docstring variable rename --- pytensor/tensor/basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index d31b277cd8..041aca4a94 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -781,7 +781,7 @@ def switch(cond, ift, iff): with pm.Model(): x = pm.Categorical('x', np.array([0.1, 0.9])) - y_ = pm.Bernoulli('y_', p=switch(x, 0.9, 0.1), shape=10) + y = pm.Bernoulli('y', p=switch(x, 0.9, 0.1), shape=10) """ where = switch From b97633181ebf64826fe6c305fa51eab6570f957e Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 08:41:05 -0500 Subject: [PATCH 4/9] Silly ruff formatting pickiness --- pytensor/tensor/basic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 041aca4a94..07776528d0 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -784,6 +784,7 @@ def switch(cond, ift, iff): y = pm.Bernoulli('y', p=switch(x, 0.9, 0.1), shape=10) """ + where = switch From b14e5646b5e3830ae2de130b194217e296ea6f6e Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 09:38:03 -0500 Subject: [PATCH 5/9] Docstring namespace fixes --- pytensor/tensor/basic.py | 4 +++- pytensor/tensor/math.py | 12 ++++++++++-- pytensor/tensor/special.py | 3 +-- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 07776528d0..758b00803f 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -779,9 +779,11 @@ def switch(cond, ift, iff): .. code:: python + import pymc as pm + with pm.Model(): x = pm.Categorical('x', np.array([0.1, 0.9])) - y = pm.Bernoulli('y', p=switch(x, 0.9, 0.1), shape=10) + y = pm.Bernoulli('y', p=pm.math.switch(x, 0.9, 0.1), shape=10) """ diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index a5668d14d3..44411acaac 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -1951,12 +1951,11 @@ def dot(l, r): .. code:: python import pymc as pm - import pytensor.tensor as pt with pm.Model() as model: x = pm.Normal('x', mu=0, sd=1, shape=2) y = pm.Normal('y', mu=0, sd=1, shape=2) - z = pt.dot(x, y) + z = pt.math.dot(x, y) """ if not isinstance(l, Variable): @@ -2676,6 +2675,15 @@ def prod( If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the original tensor. + PyMC usage example: + + .. code-block:: python + + import pymc as pm + + with pm.Model() as model: + n = pm.Poisson('n', 1, shape=(2, 3)) + prod_n = pm.Deterministic('prod_n', pm.math.prod(n, axis=0)) """ diff --git a/pytensor/tensor/special.py b/pytensor/tensor/special.py index b743c1dadb..b82988fb77 100644 --- a/pytensor/tensor/special.py +++ b/pytensor/tensor/special.py @@ -507,11 +507,10 @@ def softmax(c, axis=None): classification problem. Here is an example:: import pymc as pm - import pytensor.tensor as pt with pm.Model() as model: weights = pm.Gamma('weights', 1, 1, shape=3) - softmax_prob = pt.softmax(weights) + softmax_prob = pm.math.softmax(weights) outcome = pm.Categorical('outcome', p=softmax_prob) """ From 9c46bce69d05f5ec5f210216e0fc0933cbaba581 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 09:38:54 -0500 Subject: [PATCH 6/9] Docstring namespace fixes --- pytensor/tensor/math.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 44411acaac..c79a9552a1 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -2675,8 +2675,8 @@ def prod( If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the original tensor. - PyMC usage example: - + Example + ------- .. code-block:: python import pymc as pm From 1ca4985f306fcb0ffd95688d5b0ffdb5f5b7fe35 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 09:42:20 -0500 Subject: [PATCH 7/9] Ignore stupid ruff advice --- pytensor/tensor/math.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index c79a9552a1..ec3a20954d 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -2675,6 +2675,7 @@ def prod( If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the original tensor. + # noqa W293 Example ------- .. code-block:: python From f525c8d10299bb11bbf7c8863a1638f81e4112fe Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 09:53:03 -0500 Subject: [PATCH 8/9] cumsum usage example --- pytensor/tensor/extra_ops.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/extra_ops.py b/pytensor/tensor/extra_ops.py index 06a82744b2..8b04837200 100644 --- a/pytensor/tensor/extra_ops.py +++ b/pytensor/tensor/extra_ops.py @@ -410,9 +410,18 @@ def cumsum(x, axis=None): axis The axis along which the cumulative sum is computed. The default (None) is to compute the cumsum over the flattened array. + # noqa W293 + Example + ------- + Usage in PyMC: + .. code-block:: python - .. versionadded:: 0.7 + with pm.Model() as model: + x0 = pm.Normal('x0') + x = pm.Normal('x', mu=0, sd=1, shape=10) + # Gaussian random walk + grw = pm.Deterministic('grw', x0 + pm.math.cumsum(x)) """ return CumOp(axis=axis, mode="add")(x) From 8b62e81ab314fc04b847dd0f12c17abe7b871512 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 17 Jun 2024 10:00:19 -0500 Subject: [PATCH 9/9] cumprod example --- pytensor/tensor/extra_ops.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/pytensor/tensor/extra_ops.py b/pytensor/tensor/extra_ops.py index 8b04837200..4fc7f385b9 100644 --- a/pytensor/tensor/extra_ops.py +++ b/pytensor/tensor/extra_ops.py @@ -439,7 +439,19 @@ def cumprod(x, axis=None): axis The axis along which the cumulative product is computed. The default (None) is to compute the `cumprod` over the flattened array. + # noqa W293 + Example + ------- + Usage in PyMC: + .. code-block:: python + + import pymc as pm + + with pm.Model() as model: + x = pm.Normal('x', shape=(10, 3)) + # Product of x + prod_x = pm.Deterministic('prod_x', pm.math.cumprod(x, axis=0)) .. versionadded:: 0.7