This is now possible:

import pymc as pm
import numpy as np

def maxwell_dist(scale, size):
  return pm.math.sqrt(pm.ChiSquared.dist(nu=3, size=size)) * scale

scale = 5.0
x = pm.CustomDist.dist(scale, dist=maxwell_dist)

trunc_x = pm.Truncated.dist(x, lower=0, upper=2, size=(1000,))
assert np.all(pm.draw(trunc_x) < 2)

trunc_x = pm.Truncated.dist(x, lower=0, upper=5, size=())
assert pm.logp(trunc_x, 3.0).eval() > pm.logp(x, 3.0).eval()

This required cleaning up the interface of SymbolicRandomVariables (mainly circumventing pymc-devs/pytensor#473) so that we can safely "box" the base RVs in the inner OpFromGraph (i.e., recreate them with new shared inputs).

This challenge is very specific to Truncated which needs to "resample" the base RV for the rejection based algorithm.
No other SymbolicRandomVariable needs to do this, and they have avoided the need to box the base RVs by simply resizing them to the total size and using the resized RVs as explicit inputs to the inner graph.

For instance, Mixture will resize the component RVs to the "total size" and then scholastically index them based on its internal Categorical RV. ZeroSumNormal will create Normals as inputs and simply subtraction the mean.

Such an approach, however, makes it tricky for Truncated to know exactly what constitutes the "true" inputs of underlying SymbolicRandomVariables, and for this reason it rejected and still rejects arbitrary SymbolicRandomVariables. The exception, are the SymbolicRandomVariables created via CustomDist because for those are already "pre-boxed" in a sense. We know the relevant graph must start at dist.owner.inputs. Now that our class can safely manage and replace shared RNGs inputs, we can allow Truncated to handle such RVs, even if they require a bunch of shared RNGs.

Related to #6905 (comment)

TODO

• This is marked as draft because it includes commit from Fix rejection-based truncation of scalar variables #6923
• Failing due to Fix RandomVariable static type shape bug pytensor#475