Derive Categorical from Argmin/Argmax

pymc/distributions/custom.py

@@ -40,7 +40,6 @@
 from pymc.distributions.shape_utils import _change_dist_size, rv_size_is_none
 from pymc.exceptions import BlockModelAccessError
 from pymc.logprob.abstract import _logcdf, _logprob
-from pymc.model.core import new_or_existing_block_model_access
 from pymc.pytensorf import collect_default_updates
 
 
@@ -265,6 +264,8 @@ def rv_op(
         class_name: str,
         rng=None,
     ):
+        from pymc.model.core import new_or_existing_block_model_access
+
         size = normalize_size_param(size)
         # If it's NoneConst, just use that as the dummy
         dummy_size_param = size.type() if isinstance(size, TensorVariable) else size
@@ -827,11 +828,13 @@ def check_valid_dist_random(cls, dist, random, dist_params):
 
     @classmethod
     def is_symbolic_random(self, random, dist_params):
+        from pymc.model.core import new_or_existing_block_model_access
+
         if random is None:
             return False
         # Try calling random with symbolic inputs
+        size = normalize_size_param(None)
         try:
-            size = normalize_size_param(None)
             with new_or_existing_block_model_access(
                 error_msg_on_access="Model variables cannot be created in the random function. Use the `.dist` API to create such variables."
             ):

pymc/distributions/distribution.py

@@ -54,7 +54,6 @@
 from pymc.logprob.abstract import MeasurableOp, _icdf, _logccdf, _logcdf, _logprob
 from pymc.logprob.basic import logp
 from pymc.logprob.rewriting import logprob_rewrites_db
-from pymc.printing import str_for_dist
 from pymc.pytensorf import (
     collect_default_updates_inner_fgraph,
     constant_fold,
@@ -524,9 +523,9 @@ def __new__(
         rv : TensorVariable
             The created random variable tensor, registered in the Model.
         """
-        try:
-            from pymc.model import Model
+        from pymc.model.core import Model
 
+        try:
             model = Model.get_context()
         except TypeError:
             raise TypeError(
@@ -565,6 +564,8 @@ def __new__(
         )
 
         # add in pretty-printing support
+        from pymc.printing import str_for_dist
+
         rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
         rv_out._repr_latex_ = types.MethodType(
             functools.partial(str_for_dist, formatting="latex"), rv_out

pymc/distributions/shape_utils.py

@@ -30,7 +30,6 @@
 from pytensor.tensor.type_other import NoneTypeT
 from pytensor.tensor.variable import TensorVariable
 
-from pymc.model import modelcontext
 from pymc.pytensorf import convert_observed_data, resolve_shapes
 
 __all__ = [
@@ -386,6 +385,8 @@ def get_support_shape(
         inferred_support_shape = [shape[i] - support_shape_offset[i] for i in range(-ndim_supp, 0)]
 
     if inferred_support_shape is None and dims is not None:
+        from pymc.model import modelcontext
+
         dims = convert_dims(dims)
         assert isinstance(dims, tuple)
         if len(dims) < ndim_supp:

pymc/logprob/order.py

@@ -35,14 +35,25 @@
 #   SOFTWARE.
 from typing import cast
 
+import numpy as np
 import pytensor.tensor as pt
 
+from numpy.lib.array_utils import normalize_axis_tuple
+from pytensor.graph import ancestors
 from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
-from pytensor.tensor.math import Max
+from pytensor.scalar import Add, Mul
+from pytensor.tensor import get_underlying_scalar_constant_value
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.exceptions import NotScalarConstantError
+from pytensor.tensor.math import Argmax, Max, variadic_add, variadic_mul
+from pytensor.tensor.random.basic import ExponentialRV, GumbelRV
+from pytensor.tensor.rewriting.basic import broadcasted_by
+from pytensor.tensor.type_other import NoneTypeT
 from pytensor.tensor.variable import TensorVariable
 
+from pymc.distributions.continuous import WeibullBetaRV
 from pymc.logprob.abstract import (
     MeasurableElemwise,
     MeasurableOp,
@@ -122,6 +133,7 @@ def find_measurable_max(fgraph: FunctionGraph, node: Apply) -> list[TensorVariab
     "find_measurable_max",
     find_measurable_max,
     "basic",
+    "order",
     "max",
 )
 
@@ -158,3 +170,176 @@ def max_logprob_discrete(op, values, base_rv, **kwargs):
 
     n = pt.prod(base_rv_shape)
     return logdiffexp(n * logcdf, n * logcdf_prev)
+
+
+@node_rewriter([ExponentialRV, GumbelRV, WeibullBetaRV])
+def lift_loc_scale(fgraph, node):
+    """Rewrite rv(loc, scale) * s + l as rv(loc * s + l, scale * s), when they admit such parametrization.
+
+    Currently, this rewrite targets just RVs accepted by categorical_from_argmax, but can be generalized beyond this use case.
+    """
+    rv = node.out
+    clients = fgraph.clients[rv]
+    if len(clients) != 1:
+        return None
+    [(add_mul_node, idx)] = clients
+
+    if not (
+        isinstance(add_mul_node.op, Elemwise) and isinstance(add_mul_node.op.scalar_op, Add | Mul)
+    ):
+        return None
+
+    match node.op:
+        case ExponentialRV():
+            loc = 0
+            rng, size, scale = node.inputs
+        case WeibullBetaRV():
+            loc = 0
+            rng, size, shape, scale = node.inputs
+        case GumbelRV():
+            rng, size, loc, scale = node.inputs
+        case _:
+            raise NotImplementedError(f"Unexpected op {node.op}")
+
+    negative_scale = False
+    if isinstance(add_mul_node.op.scalar_op, Add):
+        if not isinstance(node.op, GumbelRV):
+            # Only Gumbel allows lifting loc
+            return None
+        loc += variadic_add(*(t for i, t in enumerate(add_mul_node.inputs) if i != idx))
+    else:
+        extra_scale = variadic_mul(*(t for i, t in enumerate(add_mul_node.inputs) if i != idx))
+        # The rewrite is only valid if the scale non-negative.
+        try:
+            [extra_scale_const] = constant_fold([extra_scale])
+        except NotScalarConstantError:
+            # Not-constant scale. Get out
+            # TODO: Allow more cases when we have standard machinery to infer sign of symbolic operations
+            return None
+        unique_sign = np.unique(np.sign(extra_scale_const))
+        if unique_sign.size == 2:
+            # There is mixed sign in the scale (or zero). Get out
+            return None
+        negative_scale = unique_sign == -1
+        if negative_scale:
+            if (extra_scale_const == -1).all():
+                # There is no scale to lift, it's just a negated rv (happens in argmin(rv))
+                return None
+            # Scale is homogenously negative, make it positive, and return rv * -1 later so other rewrites can handle it
+            extra_scale *= -1
+        loc *= extra_scale
+        scale *= extra_scale
+
+    # We can't lift the argument if either
+    # 1. the argument it's broadcasting the RV
+    # 2. the RV shows up in the lifted args (as in rv + rv or rv * rv)
+    # We check with loc as that is altered in both branches
+    if broadcasted_by(rv, loc) or rv in ancestors([loc]):
+        return None
+
+    match node.op:
+        case ExponentialRV():
+            lifted_rv = node.op.make_node(rng, size, scale).out
+        case WeibullBetaRV():
+            # WeibullBetaRV is a SymbolicRandomVariable, we can't simply pass arguments of a different type
+            lifted_rv = WeibullBetaRV.rv_op(shape, scale, rng=rng, size=size)
+        case GumbelRV():
+            lifted_rv = node.op.make_node(rng, size, loc, scale).out
+
+    if negative_scale:
+        lifted_rv *= -1
+
+    return {add_mul_node.out: lifted_rv}
+
+
+@node_rewriter([Argmax])
+def categorical_from_argmax(fgraph, node):
+    """Convert closed from argmax/argmin to equivalent categorical."""
+
+    def is_minus_1(x):
+        try:
+            return get_underlying_scalar_constant_value(x, max_recur=3) == -1
+        except NotScalarConstantError:
+            return False
+
+    [base_var] = node.inputs
+    base_node = base_var.owner
+
+    if base_node is None:
+        return None
+    if not filter_measurable_variables([base_var]):
+        return None
+
+    argmax_axes = node.op.axis
+    if argmax_axes is None:
+        argmax_axes = tuple(range(base_var.ndim))
+    else:
+        argmax_axes = normalize_axis_tuple(argmax_axes, base_var.ndim)
+
+    probs = None
+
+    if isinstance(base_node.op, GumbelRV):
+        # argmax(gumbel(loc, scale)) -> categorical(exp(loc / scale))
+        rng, size, loc, scale = base_node.inputs
+
+        # gumbel scale has to be constant across the argmax axes
+        if not all(b for i, b in enumerate(scale.type.broadcastable) if i in argmax_axes):
+            return None
+
+        # Dividing by scale also broadcasts probs to the size implied by (loc, scale)
+        probs = pt.exp(loc / scale)
+
+    # Check if we have an Argmin
+    # Argmin is internally represented as Argmax(-x), and -x is canonicalized as x * -1
+    elif (
+        len(base_node.inputs) == 2
+        and isinstance(base_node.op, Elemwise)
+        and isinstance(base_node.op.scalar_op, Mul)
+        and ((right_neg := is_minus_1(base_node.inputs[1])) or is_minus_1(base_node.inputs[0]))
+    ):
+        base_var = base_node.inputs[0 if right_neg else 1]
+        base_node = base_var.owner
+        if base_node is None:
+            return None
+
+        if isinstance(base_node.op, ExponentialRV):
+            # argmin(exponential(rate)) -> categorical(rate)
+            rng, size, scale = base_var.owner.inputs
+            probs = 1 / scale
+
+        elif isinstance(base_node.op, WeibullBetaRV):
+            # argmin(weibull(shape, scale)) -> categorical(scale ** -shape / Σ(scale ** -shape))
+            rng, size, shape, scale = base_node.inputs
+
+            # weibull shape has to be constant across the argmin axes
+            if not all(b for i, b in enumerate(shape.type.broadcastable) if i in argmax_axes):
+                return None
+
+            probs = scale**-shape
+
+    if probs is None:
+        return None
+
+    if not isinstance(size.type, NoneTypeT):
+        # Make probs explicit to facilitate logic below
+        probs = pt.broadcast_to(probs, size)
+
+    # Join axes probs at the last axis (core axis of Categorical)
+    n_axes = len(argmax_axes)
+    probs = pt.moveaxis(probs, argmax_axes, tuple(range(-n_axes, 0)))
+    probs = pt.join_dims(probs, -n_axes, n_axes)
+
+    # Normalize probs and create categorical
+    probs /= probs.sum(-1, keepdims=True)
+    return [pt.random.categorical(probs, rng=rng, size=None)]
+
+
+measurable_ir_rewrites_db.register("lift_loc_scale", lift_loc_scale, "basic", "lift_rv_args")
+
+measurable_ir_rewrites_db.register(
+    "categorical_from_argmax",
+    categorical_from_argmax,
+    "basic",
+    "order",
+    "argmax",
+)

pymc/logprob/transforms.py

@@ -103,6 +103,8 @@
     sub,
     tanh,
     true_div,
+    variadic_add,
+    variadic_mul,
 )
 from pytensor.tensor.variable import TensorVariable
 
@@ -518,12 +520,12 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Apply) -> list[Varia
         transform_inputs = (measurable_input, power)
         transform = PowerTransform(power=power)
     elif isinstance(scalar_op, Add):
-        transform_inputs = (measurable_input, pt.add(*other_inputs))
+        transform_inputs = (measurable_input, variadic_add(*other_inputs))
         transform = LocTransform(
             transform_args_fn=lambda *inputs: inputs[-1],
         )
     elif isinstance(scalar_op, Mul):
-        transform_inputs = (measurable_input, pt.mul(*other_inputs))
+        transform_inputs = (measurable_input, variadic_mul(*other_inputs))
         transform = ScaleTransform(
             transform_args_fn=lambda *inputs: inputs[-1],
         )