Add PyMC code examples to math function docstrings

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
 

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):

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)