I found a trivial workaround.
When using pmjax.sample_numpyro_nuts() for probabilistic modeling I realized that this function is designed for continuous distributions and doesn’t support discrete ones directly. One workaround for this limitation could be to substitute the discreteUniform distribution with a Uniform distribution. However, this substitution results in a continuous output, which may not be suitable for functions like adstock_transform that expect a discrete input. To address this, we can convert the continuous output back to a discrete value by using max_lag = int(max_lag.eval()). This snippet effectively casts the floating-point number into an integer, making it compatible with the adstock_transform function. This answer is open to feedback.