Tweak pm.Data to create or auto-replace by TensorConstant

[michaelosthege]

Background

pm.Data is often used to track data in the model, making it end up in the InferenceData and so on.

Right now this always creates shared variables, but in the vast majority of applications one doesn't actually use the shared variable features.

Generally there are three (?) flavors of data variables:

• constant (could be a TensorConstant instead)
• shared with constant shape
• shared with flexible shape (constant ndim)

AFAIK Aesara does not distinguish between shared variables that can, or cannot be resized.

Motivation

But the difference between TensorConstant and SharedVariable is a thing, and sometimes also costly.
For example see #3818

Potential solution

• Additional kwargs to pm.Data through which one can constrain the flexibility (creating TensorConstant instead).
• Automatically replacing all SharedVariables in the graph with TensorConstant before starting MCMC.

Versions and main components

• PyMC/PyMC3 Version: main

[ricardoV94]

Automatically replacing all SharedVariables in the graph with TensorConstant before starting MCMC.

I implemented this for other reasons here:

pymc/pymc/sampling_jax.py

Line 46 in 9e7e8aa

def replace_shared_variables(graph: List[TensorVariable]) -> List[TensorVariable]:

Should be possible. Perhaps a more flexible version that just skips Variables with default updates instead of raising an Error would be fine.