Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

MergeOptimization tries to merge nodes with different static shapes #955

Closed
ricardoV94 opened this issue May 11, 2022 · 2 comments · Fixed by #915
Closed

MergeOptimization tries to merge nodes with different static shapes #955

ricardoV94 opened this issue May 11, 2022 · 2 comments · Fixed by #915
Assignees
@ricardoV94
Labels
bug Something isn't working graph rewriting important shape inference

Comments

@ricardoV94
Copy link
Contributor

ricardoV94 commented May 11, 2022 •
edited

Still need to investigate this one a bit further, showed up in pymc-devs/pymc#5757

import aesara
import aesara.tensor as at
import numpy as np

w = aesara.shared(np.random.rand(20))
u = aesara.shared(np.random.rand(20))

z0 = at.arange(0, 20).astype("float64")
z0_ = z0.dimshuffle("x", 0)
hwz = at.tanh(z0_.dot(w))
z1 = z0_ + at.outer(hwz, u)
z1 = z1.flatten()
aesara.config.on_opt_error = "raise"
at.jacobian(z1, z0).eval()
Traceback 
<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> The type of the replacement (TensorType(float64, (1, None))) must be compatible with the type of the original Variable (TensorType(float64, (1, 20))). MergeOptimizer
ERROR (aesara.graph.opt): SeqOptimizer apply MergeOptimizer
ERROR (aesara.graph.opt): Traceback:
ERROR (aesara.graph.opt): Traceback (most recent call last):
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 280, in apply
    sub_prof = optimizer.optimize(fgraph)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 103, in optimize
    ret = self.apply(fgraph, *args, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 746, in apply
    fgraph.replace_all_validate(pairs, reason="MergeOptimizer")
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/features.py", line 553, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/fg.py", line 511, in replace
    self.change_node_input(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/fg.py", line 428, in change_node_input
    raise TypeError(
TypeError: The type of the replacement (TensorType(float64, (1, None))) must be compatible with the type of the original Variable (TensorType(float64, (1, 20))).
Traceback (most recent call last):
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/link/vm.py", line 1112, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/scan/op.py", line 1411, in make_thunk
    isinstance(out, TensorVariable) for out in self.fn.maker.fgraph.outputs
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/scan/op.py", line 1349, in fn
    self._fn = function(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/__init__.py", line 317, in function
    fn = pfunc(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/pfunc.py", line 363, in pfunc
    return orig_function(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/types.py", line 1725, in orig_function
    m = Maker(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/types.py", line 1464, in __init__
    optimizer_profile = optimizer(fgraph)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 112, in __call__
    return self.optimize(fgraph)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 103, in optimize
    ret = self.apply(fgraph, *args, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 291, in apply
    self.failure_callback(e, self, optimizer)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 226, in warn
    raise exc
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 280, in apply
    sub_prof = optimizer.optimize(fgraph)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 103, in optimize
    ret = self.apply(fgraph, *args, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 746, in apply
    fgraph.replace_all_validate(pairs, reason="MergeOptimizer")
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/features.py", line 553, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/fg.py", line 511, in replace
    self.change_node_input(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/fg.py", line 428, in change_node_input
    raise TypeError(
TypeError: The type of the replacement (TensorType(float64, (1, None))) must be compatible with the type of the original Variable (TensorType(float64, (1, 20))).
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3418, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-37-615f4005dd2e>", line 15, in <module>
    at.jacobian(z1, z0).eval()
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/basic.py", line 567, in eval
    self._fn_cache[inputs] = function(inputs, self)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/__init__.py", line 317, in function
    fn = pfunc(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/pfunc.py", line 363, in pfunc
    return orig_function(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/types.py", line 1736, in orig_function
    fn = m.create(defaults)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/types.py", line 1631, in create
    _fn, _i, _o = self.linker.make_thunk(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/link/basic.py", line 254, in make_thunk
    return self.make_all(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/link/vm.py", line 1121, in make_all
    raise_with_op(fgraph, node)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/link/utils.py", line 538, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/link/vm.py", line 1112, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/scan/op.py", line 1411, in make_thunk
    isinstance(out, TensorVariable) for out in self.fn.maker.fgraph.outputs
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/scan/op.py", line 1349, in fn
    self._fn = function(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/__init__.py", line 317, in function
    fn = pfunc(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/pfunc.py", line 363, in pfunc
    return orig_function(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/types.py", line 1725, in orig_function
    m = Maker(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/compile/function/types.py", line 1464, in __init__
    optimizer_profile = optimizer(fgraph)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 112, in __call__
    return self.optimize(fgraph)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 103, in optimize
    ret = self.apply(fgraph, *args, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 291, in apply
    self.failure_callback(e, self, optimizer)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 226, in warn
    raise exc
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 280, in apply
    sub_prof = optimizer.optimize(fgraph)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 103, in optimize
    ret = self.apply(fgraph, *args, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/opt.py", line 746, in apply
    fgraph.replace_all_validate(pairs, reason="MergeOptimizer")
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/features.py", line 553, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/fg.py", line 511, in replace
    self.change_node_input(
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/aesara/graph/fg.py", line 428, in change_node_input
    raise TypeError(
TypeError: The type of the replacement (TensorType(float64, (1, None))) must be compatible with the type of the original Variable (TensorType(float64, (1, 20))).
Apply node that caused the error: for{cpu,scan_fn}(TensorConstant{20}, TensorConstant{[ 0  1  2 .. 17 18 19]}, TensorConstant{20}, Reshape{1}.0, InplaceDimShuffle{x,0}.0, Elemwise{Composite{(i0 - sqr(i1))}}[(0, 1)].0, InplaceDimShuffle{0,x}.0)
Toposort index: 8
Inputs types: [TensorType(int64, ()), TensorType(int64, (20,)), TensorType(int64, ()), TensorType(float64, (None,)), TensorType(float64, (1, None)), TensorType(float64, (1,)), TensorType(float64, (None, 1))]
HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most Aesara optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the Aesara flag 'optimizer=fast_compile'. If that does not work, Aesara optimizations can be disabled with 'optimizer=None'.
HINT: Use the Aesara flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
@ricardoV94 ricardoV94 changed the title Optimization error due to underspecified static shape MergeOptimization fails when trying to merge nodes with different static shapes May 11, 2022
@ricardoV94 ricardoV94 changed the title MergeOptimization fails when trying to merge nodes with different static shapes MergeOptimization tries to merge nodes with different static shapes May 11, 2022
@ricardoV94
Copy link
Contributor Author

ricardoV94 commented May 11, 2022 •
edited

Somehow there are two Dimshuffles with the same inputs, but different output static shape, not sure how this can happen. Some inplace change to the apply node?

Adding a breakpoint after this if statement triggers it:

aesara/aesara/graph/opt.py

Lines 637 to 639 in d11f303

if (node, candidate) in self.blacklist:
# They were already tried, and there was an error
continue 

if node.outputs[0].type != candidate.outputs[0].type:
    breakpoint()

@brandonwillard
Copy link
Member

It looks like we need to apply the Type ordering (e.g. Type.is_super) to the merge candidates in MergeOptimizer.apply.

@ricardoV94 ricardoV94 self-assigned this May 12, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@ricardoV94 ricardoV94

Labels
bug Something isn't working graph rewriting important shape inference
Projects
None yet
Milestone
No milestone
2 participants
@brandonwillard @ricardoV94