potential bug - why are variable ids changing when used in an expression?

[kratsg]

As part of a library (https://github.com/scipp-atlas/pyhs3/), we keep track of the inputs globally based on what the user provides to be able to compile functions later, however there was some interesting behavior where the id of the variable is changing... when I would have presumed it to not have been cloned?

from pytensor.printing import debugprint
import pytensor.tensor as pt
from pytensor.graph.traversal import explicit_graph_inputs

x = pt.vector("x")[:, None]
y = pt.vector("y")[:, None]
c = pt.scalar("c")

print(f'id(x) = {id(x)}')
print(f'id(y) = {id(y)}')
print(f'id(c) = {id(c)}')

result = pt.math.exp(x*y*c)

print('debugprint(result)')
debugprint(result)
for variable in explicit_graph_inputs([result]):
    print(f'variable={variable}, id={id(variable)}, debugprint:')
    debugprint(variable)

The output of the above gives

id(x) = 4593018608
id(y) = 4598833728
id(c) = 4417004672
debugprint(result)
Exp [id A]
 └─ Mul [id B]
    ├─ Mul [id C]
    │  ├─ ExpandDims{axis=1} [id D]
    │  │  └─ x [id E]
    │  └─ ExpandDims{axis=1} [id F]
    │     └─ y [id G]
    └─ ExpandDims{axes=[0, 1]} [id H]
       └─ c [id I]
variable=c, id=4417004672, debugprint:
c [id A]
variable=y, id=4598831376, debugprint:
y [id A]
variable=x, id=4344565536, debugprint:
x [id A]

and I noticed the x and the y have their ids changed, but the c (scalar) does not. Is this intentional or not? I'm a bit confused about what happens under the hood for what seems to be slightly different behavior here? Presumably, the ExpandDims is getting stripped through explicit_graph_inputs - or that the input should be treated as x and y as vectors, and part of the computational graph will be to expand them?

[kratsg]

For further context, I am trying to reproduce or MWE this problematic code below:

>>> function([y_var, c_var, x_var], dist_expr, mode='DebugMode')
Traceback (most recent call last):
  File "<console>", line 1, in <module>
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/compile/function/__init__.py", line 332, in function
    fn = pfunc(
         ^^^^^^
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/compile/function/pfunc.py", line 464, in pfunc
    return orig_function(
           ^^^^^^^^^^^^^^
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 1821, in orig_function
    m = Maker(
        ^^^^^^
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/compile/debugmode.py", line 2266, in function_maker
    return _Maker(i, o, self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/compile/debugmode.py", line 2038, in __init__
    fgraph, additional_outputs, equivalence_tracker = _optcheck_fgraph(
                                                      ^^^^^^^^^^^^^^^^^
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/compile/debugmode.py", line 436, in _optcheck_fgraph
    fgraph, updates = std_fgraph(
                      ^^^^^^^^^^^
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 276, in std_fgraph
    fgraph = FunctionGraph(
             ^^^^^^^^^^^^^^
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/graph/fg.py", line 166, in __init__
    self.add_output(output, reason="init")
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/graph/fg.py", line 176, in add_output
    self.import_var(var, reason=reason, import_missing=import_missing)
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/graph/fg.py", line 316, in import_var
    self.import_node(apply, reason=reason, import_missing=import_missing)
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pytensor/graph/fg.py", line 382, in import_node
    raise MissingInputError(error_msg, variable=var)
pytensor.graph.utils.MissingInputError: Input 0 (y) of the graph (indices start from 0), used to compute ExpandDims{axis=1}(y), was not provided and not given a value. Use the PyTensor flag exception_verbosity='high', for more information on this error.

Backtrace when that variable is created:

  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pluggy/_manager.py", line 120, in _hookexec
    return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/pluggy/_callers.py", line 121, in _multicall
    res = hook_impl.function(*args)
  File "/Users/kratsg/pyhs3/.pixi/envs/py312-jax/lib/python3.12/site-packages/_pytest/python.py", line 166, in pytest_pyfunc_call
    result = testfunction(**testargs)
  File "/Users/kratsg/pyhs3/tests/test_normalization_integration.py", line 87, in test_model_with_two_observables
    model = Model(
  File "/Users/kratsg/pyhs3/src/pyhs3/model.py", line 113, in __init__
    self._build_dependency_graph(functions, distributions, progress)
  File "/Users/kratsg/pyhs3/src/pyhs3/model.py", line 405, in _build_dependency_graph
    self.parameters[node_name] = self._build_parameter_node(
  File "/Users/kratsg/pyhs3/src/pyhs3/model.py", line 305, in _build_parameter_node
    tensor = create_bounded_tensor(node_name, domain_bounds, param_kind)
  File "/Users/kratsg/pyhs3/src/pyhs3/model.py", line 788, in create_bounded_tensor
    tensor = kind(name)

where this version works

>>> function([y_var.owner.inputs[0], c_var, x_var.owner.inputs[0]], dist_expr)
<pytensor.compile.function.types.Function object at 0x13907f860>

however, I'm likely running into problems where my inputs need to be the variable, not the expanded-dimensions-variable when I compile the function?

[kratsg]

Ahh, I think I'm understanding and I hit a red-herring when I'm trying to debug. I should keep track of the variable, not the expanded-dimension one, when I'm trying to build these graphs and use clone_replace. Here's why:

>>> debugprint(dist_expr.owner.inputs[0].owner.inputs[0])
Exp [id A]
 └─ Mul [id B]
    ├─ Mul [id C]
    │  ├─ ExpandDims{axes=[0, 1]} [id D]
    │  │  └─ c [id E]
    │  └─ ExpandDims{axis=1} [id F] 'x'
    │     └─ x [id G]
    └─ ExpandDims{axis=1} [id H] 'y'
       └─ y [id I]

>>> function([x_var, y_var, c_var], dist_expr.owner.inputs[0].owner.inputs[0])
<pytensor.compile.function.types.Function object at 0x1385dfd40>

This works absolutely fine (a subgraph here). As part of the code, we are deriving a normalization via Gauss-Legendre approximation, and this is done by doing a clone-replace like so:

_GL_ORDER = 64
_GL_NODES, _GL_WEIGHTS = np.polynomial.legendre.leggauss(_GL_ORDER)
_GL_NODES_T = pt.constant(_GL_NODES)
_GL_WEIGHTS_T = pt.constant(_GL_WEIGHTS)
...

    half_width = (upper - lower) / 2.0
    midpoint = (upper + lower) / 2.0

    # evaluation points: (64,) for 1-D variables, (64, 1) for 2-D observables (N, 1)
    x_points = half_width * _GL_NODES_T + midpoint

    # Observable is (N, 1); replace with (64, 1) quadrature points.
    # f_vals is then (64, 1) or scalar.  Use a (64, 1) weight column so
    # broadcasting is unambiguous: (64,1)*(64,1) → (64,1), sum axis 0 → (1,).
    x_points = x_points[:, None]
    f_vals = clone_replace(expression, replace=[(variable, x_points)])
    integral = half_width * pt.squeeze(  # type: ignore[no-untyped-call]
        pt.sum(_GL_WEIGHTS_T[:, None] * f_vals, axis=0)  # type: ignore[no-untyped-call]
    )

and this ends up with the following graph:

>>> debugprint(dist_expr)
Mul [id A] 'test_dist'
 ├─ True_div [id B]
 │  ├─ Exp [id C]
 │  │  └─ Mul [id D]
 │  │     ├─ Mul [id E]
 │  │     │  ├─ ExpandDims{axes=[0, 1]} [id F]
 │  │     │  │  └─ c [id G]
 │  │     │  └─ ExpandDims{axis=1} [id H] 'x'
 │  │     │     └─ x [id I]
 │  │     └─ ExpandDims{axis=1} [id J] 'y'
 │  │        └─ y [id K]
 │  └─ ExpandDims{axes=[0, 1]} [id L]
 │     └─ Mul [id M]
 │        ├─ True_div [id N]
 │        │  ├─ Sub [id O]
 │        │  │  ├─ 5.0 [id P]
 │        │  │  └─ -5.0 [id Q]
 │        │  └─ 2.0 [id R]
 │        └─ Squeeze{axis=0} [id S]
 │           └─ Sum{axis=0} [id T]
 │              └─ Mul [id U]
 │                 ├─ ExpandDims{axis=1} [id V]
 │                 │  └─ [0.0017832 ... .00178328] [id W]
 │                 └─ ExpandDims{axes=[0, 1]} [id X]
 │                    └─ Mul [id Y]
 │                       ├─ True_div [id Z]
 │                       │  ├─ Sub [id BA]
 │                       │  │  ├─ 10.0 [id BB]
 │                       │  │  └─ 0.0 [id BC]
 │                       │  └─ 2.0 [id BD]
 │                       └─ Squeeze{axis=0} [id BE]
 │                          └─ Sum{axis=0} [id BF]
 │                             └─ Mul [id BG]
 │                                ├─ ExpandDims{axis=1} [id BH]
 │                                │  └─ [0.0017832 ... .00178328] [id W]
 │                                └─ Exp [id BI]
 │                                   └─ Mul [id BJ]
 │                                      ├─ Mul [id BK]
 │                                      │  ├─ ExpandDims{axes=[0, 1]} [id BL]
 │                                      │  │  └─ c [id G]
 │                                      │  └─ ExpandDims{axis=1} [id BM]
 │                                      │     └─ Add [id BN]
 │                                      │        ├─ Mul [id BO]
 │                                      │        │  ├─ ExpandDims{axis=0} [id BP]
 │                                      │        │  │  └─ True_div [id BQ]
 │                                      │        │  │     ├─ Sub [id BR]
 │                                      │        │  │     │  ├─ 10.0 [id BB]
 │                                      │        │  │     │  └─ 0.0 [id BC]
 │                                      │        │  │     └─ 2.0 [id BD]
 │                                      │        │  └─ [-0.999305 ... .99930504] [id BS]
 │                                      │        └─ ExpandDims{axis=0} [id BT]
 │                                      │           └─ True_div [id BU]
 │                                      │              ├─ Add [id BV]
 │                                      │              │  ├─ 10.0 [id BB]
 │                                      │              │  └─ 0.0 [id BC]
 │                                      │              └─ 2.0 [id BW]
 │                                      └─ ExpandDims{axis=1} [id BX] 'y'
 │                                         └─ y [id K]
 └─ ExpandDims{axes=[0, 1]} [id BY]
    └─ 1.0 [id BZ]

but the problem is I was not expecting a y in the denominator here, but the clone_replace cloned ALL other variables and created a separate y here (whoops!). Basically, I am trying to integrate multiple times.

[kratsg]

turns out graph_replace is what I needed, not clone_replace.