generalised expotential function added

pymc/logprob/transforms.py

@@ -556,19 +556,29 @@ def measurable_special_log_to_log(fgraph, node):
         return [pt.log(inp) / pt.log(10)]
 
 
-@node_rewriter([exp2, expm1, sigmoid])
+@node_rewriter([expm1, sigmoid])
 def measurable_special_exp_to_exp(fgraph, node):
-    """Convert log1p, log1mexp, softplus of `MeasurableVariable`s to log form."""
+    """Convert expm1, sigmoid of `MeasurableVariable`s to xp form."""
     [inp] = node.inputs
-
-    if isinstance(node.op.scalar_op, Exp2):
-        return [pt.power(2, inp)]
     if isinstance(node.op.scalar_op, Expm1):
         return [pt.exp(inp) - 1]
     if isinstance(node.op.scalar_op, Sigmoid):
         return [1 / (1 + pt.exp(-inp))]
 
 
+@node_rewriter([exp2, pow])
+def measurable_general_exp_to_exp(fgraph, node):
+    """Convert exp2 and any const^x of `MeasurableVariable`s to exp form."""
+    if len(node.inputs) > 1:
+        [const, inp] = node.inputs
+    else:
+        [inp] = node.inputs
+    if isinstance(node.op.scalar_op, Exp2):
+        return [pt.exp(pt.log(2) * inp)]
+    if isinstance(node.op.scalar_op, Pow) and isinstance(inp, pt.TensorVariable):
+        return [pt.exp(pt.log(const) * inp)]
+
+
 @node_rewriter(
     [
         exp,
@@ -721,6 +731,13 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
     "transform",
 )
 
+measurable_ir_rewrites_db.register(
+    "measurable_general_exp_to_exp",
+    measurable_general_exp_to_exp,
+    "basic",
+    "transform",
+)
+
 measurable_ir_rewrites_db.register(
     "find_measurable_transforms",
     find_measurable_transforms,