Repository: pymc-devs/pymc

Issue:
ADVI errors in numba mode for `StudentT` likelihood when `total_size` is set
### Description

```py
import pymc as pm
import numpy as np

rng = np.random.default_rng()

with pm.Model() as m:
    data = pm.Data('data', rng.normal(size=(1000, 5)))
    obs = pm.Data('obs', rng.normal(size=(1000,)))
    
    data_batch, obs_batch = pm.Minibatch(data, obs, batch_size=128)
    
    beta = pm.Normal('beta', size=(5,))
    mu = data_batch @ beta
    sigma = pm.Exponential('sigma', 1)
    
    y_hat = pm.StudentT('y_hat', mu=mu, sigma=sigma, nu=3, observed=obs_batch, total_size=1000)
    
    idata = pm.fit(n=1_000_000, 
                   compile_kwargs={'mode':'NUMBA'})
```

<details>

<summary>Full Traceback</summary>

```bash

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[1], line 18
     14 sigma = pm.Exponential('sigma', 1)
     16 y_hat = pm.StudentT('y_hat', mu=mu, sigma=sigma, nu=3, observed=obs_batch, total_size=1000)
---> 18 idata = pm.fit(n=1_000_000, 
     19                compile_kwargs={'mode':'NUMBA'})

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pymc/variational/inference.py:775, in fit(n, method, model, random_seed, start, start_sigma, inf_kwargs, **kwargs)
    773 else:
    774     raise TypeError(f"method should be one of {set(_select.keys())} or Inference instance")
--> 775 return inference.fit(n, **kwargs)

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pymc/variational/inference.py:158, in Inference.fit(self, n, score, callbacks, progressbar, progressbar_theme, **kwargs)
    156     callbacks = []
    157 score = self._maybe_score(score)
--> 158 step_func = self.objective.step_function(score=score, **kwargs)
    160 if score:
    161     state = self._iterate_with_loss(
    162         0, n, step_func, progressbar, progressbar_theme, callbacks
    163     )

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/configparser.py:44, in _ChangeFlagsDecorator.__call__.<locals>.res(*args, **kwargs)
     41 @wraps(f)
     42 def res(*args, **kwargs):
     43     with self:
---> 44         return f(*args, **kwargs)

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pymc/variational/opvi.py:405, in ObjectiveFunction.step_function(self, obj_n_mc, tf_n_mc, obj_optimizer, test_optimizer, more_obj_params, more_tf_params, more_updates, more_replacements, total_grad_norm_constraint, score, compile_kwargs, fn_kwargs)
    403 seed = self.approx.rng.randint(2**30, dtype=np.int64)
    404 if score:
--> 405     step_fn = compile([], updates.loss, updates=updates, random_seed=seed, **compile_kwargs)
    406 else:
    407     step_fn = compile([], [], updates=updates, random_seed=seed, **compile_kwargs)

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pymc/pytensorf.py:947, in compile(inputs, outputs, random_seed, mode, **kwargs)
    945 opt_qry = mode.provided_optimizer.including("random_make_inplace", check_parameter_opt)
    946 mode = Mode(linker=mode.linker, optimizer=opt_qry)
--> 947 pytensor_function = pytensor.function(
    948     inputs,
    949     outputs,
    950     updates={**rng_updates, **kwargs.pop("updates", {})},
    951     mode=mode,
    952     **kwargs,
    953 )
    954 return pytensor_function

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/compile/function/__init__.py:332, in function(inputs, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input, trust_input)
    321     fn = orig_function(
    322         inputs,
    323         outputs,
   (...)    327         trust_input=trust_input,
    328     )
    329 else:
    330     # note: pfunc will also call orig_function -- orig_function is
    331     #      a choke point that all compilation must pass through
--> 332     fn = pfunc(
    333         params=inputs,
    334         outputs=outputs,
    335         mode=mode,
    336         updates=updates,
    337         givens=givens,
    338         no_default_updates=no_default_updates,
    339         accept_inplace=accept_inplace,
    340         name=name,
    341         rebuild_strict=rebuild_strict,
    342         allow_input_downcast=allow_input_downcast,
    343         on_unused_input=on_unused_input,
    344         profile=profile,
    345         output_keys=output_keys,
    346         trust_input=trust_input,
    347     )
    348 return fn

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/compile/function/pfunc.py:466, in pfunc(params, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input, output_keys, fgraph, trust_input)
    452     profile = ProfileStats(message=profile)
    454 inputs, cloned_outputs = construct_pfunc_ins_and_outs(
    455     params,
    456     outputs,
   (...)    463     fgraph=fgraph,
    464 )
--> 466 return orig_function(
    467     inputs,
    468     cloned_outputs,
    469     mode,
    470     accept_inplace=accept_inplace,
    471     name=name,
    472     profile=profile,
    473     on_unused_input=on_unused_input,
    474     output_keys=output_keys,
    475     fgraph=fgraph,
    476     trust_input=trust_input,
    477 )

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/compile/function/types.py:1833, in orig_function(inputs, outputs, mode, accept_inplace, name, profile, on_unused_input, output_keys, fgraph, trust_input)
   1820     m = Maker(
   1821         inputs,
   1822         outputs,
   (...)   1830         trust_input=trust_input,
   1831     )
   1832     with config.change_flags(compute_test_value="off"):
-> 1833         fn = m.create(defaults)
   1834 finally:
   1835     if profile and fn:

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/compile/function/types.py:1717, in FunctionMaker.create(self, input_storage, storage_map)
   1714 start_import_time = pytensor.link.c.cmodule.import_time
   1716 with config.change_flags(traceback__limit=config.traceback__compile_limit):
-> 1717     _fn, _i, _o = self.linker.make_thunk(
   1718         input_storage=input_storage_lists, storage_map=storage_map
   1719     )
   1721 end_linker = time.perf_counter()
   1723 linker_time = end_linker - start_linker

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/basic.py:245, in LocalLinker.make_thunk(self, input_storage, output_storage, storage_map, **kwargs)
    238 def make_thunk(
    239     self,
    240     input_storage: Optional["InputStorageType"] = None,
   (...)    243     **kwargs,
    244 ) -> tuple["BasicThunkType", "InputStorageType", "OutputStorageType"]:
--> 245     return self.make_all(
    246         input_storage=input_storage,
    247         output_storage=output_storage,
    248         storage_map=storage_map,
    249     )[:3]

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/basic.py:695, in JITLinker.make_all(self, input_storage, output_storage, storage_map)
    692 for k in storage_map:
    693     compute_map[k] = [k.owner is None]
--> 695 thunks, nodes, jit_fn = self.create_jitable_thunk(
    696     compute_map, nodes, input_storage, output_storage, storage_map
    697 )
    699 [fn] = thunks
    700 fn.jit_fn = jit_fn

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/basic.py:647, in JITLinker.create_jitable_thunk(self, compute_map, order, input_storage, output_storage, storage_map)
    644 # This is a bit hackish, but we only return one of the output nodes
    645 output_nodes = [o.owner for o in self.fgraph.outputs if o.owner is not None][:1]
--> 647 converted_fgraph = self.fgraph_convert(
    648     self.fgraph,
    649     order=order,
    650     input_storage=input_storage,
    651     output_storage=output_storage,
    652     storage_map=storage_map,
    653 )
    655 thunk_inputs = self.create_thunk_inputs(storage_map)
    656 thunk_outputs = [storage_map[n] for n in self.fgraph.outputs]

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/numba/linker.py:10, in NumbaLinker.fgraph_convert(self, fgraph, **kwargs)
      7 def fgraph_convert(self, fgraph, **kwargs):
      8     from pytensor.link.numba.dispatch import numba_funcify
---> 10     return numba_funcify(fgraph, **kwargs)

File ~/mambaforge/envs/econ/lib/python3.12/functools.py:912, in singledispatch.<locals>.wrapper(*args, **kw)
    908 if not args:
    909     raise TypeError(f'{funcname} requires at least '
    910                     '1 positional argument')
--> 912 return dispatch(args[0].__class__)(*args, **kw)

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/numba/dispatch/basic.py:380, in numba_funcify_FunctionGraph(fgraph, node, fgraph_name, **kwargs)
    373 @numba_funcify.register(FunctionGraph)
    374 def numba_funcify_FunctionGraph(
    375     fgraph,
   (...)    378     **kwargs,
    379 ):
--> 380     return fgraph_to_python(
    381         fgraph,
    382         numba_funcify,
    383         type_conversion_fn=numba_typify,
    384         fgraph_name=fgraph_name,
    385         **kwargs,
    386     )

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/utils.py:736, in fgraph_to_python(fgraph, op_conversion_fn, type_conversion_fn, order, storage_map, fgraph_name, global_env, local_env, get_name_for_object, squeeze_output, unique_name, **kwargs)
    734 body_assigns = []
    735 for node in order:
--> 736     compiled_func = op_conversion_fn(
    737         node.op, node=node, storage_map=storage_map, **kwargs
    738     )
    740     # Create a local alias with a unique name
    741     local_compiled_func_name = unique_name(compiled_func)

File ~/mambaforge/envs/econ/lib/python3.12/functools.py:912, in singledispatch.<locals>.wrapper(*args, **kw)
    908 if not args:
    909     raise TypeError(f'{funcname} requires at least '
    910                     '1 positional argument')
--> 912 return dispatch(args[0].__class__)(*args, **kw)

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/numba/dispatch/random.py:401, in numba_funcify_RandomVariable(op, node, **kwargs)
    398 core_shape_len = get_vector_length(core_shape)
    399 inplace = rv_op.inplace
--> 401 core_rv_fn = numba_core_rv_funcify(rv_op, rv_node)
    402 nin = 1 + len(dist_params)  # rng + params
    403 core_op_fn = store_core_outputs(core_rv_fn, nin=nin, nout=1)

File ~/mambaforge/envs/econ/lib/python3.12/functools.py:912, in singledispatch.<locals>.wrapper(*args, **kw)
    908 if not args:
    909     raise TypeError(f'{funcname} requires at least '
    910                     '1 positional argument')
--> 912 return dispatch(args[0].__class__)(*args, **kw)

File ~/mambaforge/envs/econ/lib/python3.12/site-packages/pytensor/link/numba/dispatch/random.py:47, in numba_core_rv_funcify(op, node)
     44 @singledispatch
     45 def numba_core_rv_funcify(op: Op, node: Apply) -> Callable:
     46     """Return the core function for a random variable operation."""
---> 47     raise NotImplementedError(f"Core implementation of {op} not implemented.")

NotImplementedError: Core implementation of t_rv{"(),(),()->()"} not implemented.
```

</details>

Interestingly, it works fine if you change `total_size = None`.

Hints from issue comments:
@ricardoV94: I think we simply don't have a dispatch to StudentT in numba

Solution patch:
diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py
index 82eca936b..78eb3f7bb 100644
--- a/pymc/pytensorf.py
+++ b/pymc/pytensorf.py
@@ -979,7 +979,7 @@ def constant_fold(
     """
     fg = FunctionGraph(outputs=xs, features=[ShapeFeature()], copy_inputs=False, clone=True)
 
-    # The default rewrite_graph includes a constand_folding that is not always applied.
+    # The default rewrite_graph includes a constant_folding that is not always applied.
     # We use an unconditional constant_folding as the last pass to ensure a thorough constant folding.
     rewrite_graph(fg)
     topo_unconditional_constant_folding.apply(fg)
diff --git a/pymc/variational/minibatch_rv.py b/pymc/variational/minibatch_rv.py
index 163cec472..34d30cfa5 100644
--- a/pymc/variational/minibatch_rv.py
+++ b/pymc/variational/minibatch_rv.py
@@ -40,6 +40,9 @@ class MinibatchRandomVariable(MeasurableOp, Op):
         out = rv.type()
         return Apply(self, [rv, *total_size], [out])
 
+    def infer_shape(self, fgraph, node, shapes):
+        return [shapes[0]]
+
     def perform(self, node, inputs, output_storage):
         output_storage[0][0] = inputs[0]
 
diff --git a/pymc/variational/opvi.py b/pymc/variational/opvi.py
index 0903d2775..deedfc8d9 100644
--- a/pymc/variational/opvi.py
+++ b/pymc/variational/opvi.py
@@ -74,6 +74,7 @@ from pymc.model import modelcontext
 from pymc.pytensorf import (
     SeedSequenceSeed,
     compile,
+    constant_fold,
     find_rng_nodes,
     reseed_rngs,
 )
@@ -1105,7 +1106,10 @@ class Group(WithMemoization):
         t = self.to_flat_input(
             pt.max(
                 [
-                    get_scaling(v.owner.inputs[1:], v.shape)
+                    get_scaling(
+                        v.owner.inputs[1:],
+                        constant_fold([v.owner.inputs[0].shape], raise_not_constant=False),
+                    )
                     for v in self.group
                     if isinstance(v.owner.op, MinibatchRandomVariable)
                 ]
@@ -1272,7 +1276,10 @@ class Approximation(WithMemoization):
         t = pt.max(
             self.collect("symbolic_normalizing_constant")
             + [
-                get_scaling(obs.owner.inputs[1:], obs.shape)
+                get_scaling(
+                    obs.owner.inputs[1:],
+                    constant_fold([obs.owner.inputs[0].shape], raise_not_constant=False),
+                )
                 for obs in self.model.observed_RVs
                 if isinstance(obs.owner.op, MinibatchRandomVariable)
             ]


Test patch:
diff --git a/tests/variational/test_opvi.py b/tests/variational/test_opvi.py
index d692b3001..0f40572f7 100644
--- a/tests/variational/test_opvi.py
+++ b/tests/variational/test_opvi.py
@@ -20,6 +20,7 @@ import pytest
 
 import pymc as pm
 
+from pymc.testing import assert_no_rvs
 from pymc.variational import opvi
 from pymc.variational.approximations import (
     Empirical,
@@ -278,3 +279,18 @@ def test_logq_globals(three_var_approx):
     es = symbolic_logq.eval()
     assert e.shape == ()
     assert es.shape == (2,)
+
+
+def test_symbolic_normalizing_constant_no_rvs():
+    # Test that RVs aren't included in the graph of symbolic_normalizing_constant
+    rng = np.random.default_rng()
+
+    with pm.Model() as m:
+        obs = pm.Data("obs", rng.normal(size=(1000,)))
+        obs_batch = pm.Minibatch(obs, batch_size=128)
+        x = pm.Normal("x")  # Need at least one Free_RV in the graph
+        y_hat = pm.Flat("y_hat", observed=obs_batch, total_size=1000)
+
+        step = pm.ADVI()
+
+    assert_no_rvs(step.approx.symbolic_normalizing_constant)