Refactor infer_shape method of Ops to find output shapes using gufunc_signature
#1294
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #1294 +/- ##
==========================================
- Coverage 81.99% 81.99% -0.01%
==========================================
Files 188 188
Lines 48506 48502 -4
Branches 8672 8671 -1
==========================================
- Hits 39771 39767 -4
Misses 6583 6583
Partials 2152 2152
🚀 New features to boost your workflow:
|
|
@ricardoV94 Some Ops like |
|
That's fine we know gufunc signature won't cover all cases, the idea is to just cover the easy ones |
pytensor/tensor/nlinalg.py
Outdated
| @@ -156,7 +157,7 @@ def R_op(self, inputs, eval_points): | |||
| return [-matrix_dot(xi, ev, xi)] | |||
|
|
|||
| def infer_shape(self, fgraph, node, shapes): | |||
| return shapes | |||
| return _gufunc_to_out_shape(self.gufunc_signature, shapes) | |||
There was a problem hiding this comment.
The base class should do this if there's a signature so we can remove the infer_shape methods of these Ops altogether.
There was a problem hiding this comment.
Implemented this in my latest commit here. Now there is a infer_shape method in the base Op class, which works when there is a gufunc_signature present. But is this the right way? Now all Ops have a infer_shape method by default.
There was a problem hiding this comment.
@ricardoV94 Adding a infer_shape method to the base class seems to change function graphs. In particular, there is a test here, the second assertion fails, which checks if function graph of a function which computes the output shape of a Blockwise Op formed using MyTestOp core-op has any MyTestOp nodes in it. Ideally, no TestOp nodes should be present. But when infer_shape is added to base Op class, MyTestOp nodes appear in the function graph. This causes an assertion fail, failing that test case. How do you suggest I proceed?
Here is an example:
from tests.tensor.test_blockwise import MyTestOp
from tests.tensor.test_blockwise import Blockwise
from tests.tensor.test_blockwise import tensor
import pytensor as pt
inp = tensor(shape=(5, None, None))
test_op = MyTestOp()
op = Blockwise(test_op, signature="(m, n) -> (n, m)")
out = op(inp)
shape_fn = pt.function([inp], out.shape)
shape_fn.dprint()With infer_shape in base Op:
MakeVector{dtype='int64'} [id A] 3
├─ 5 [id B]
├─ Shape_i{1} [id C] 2
│ └─ Blockwise{MyTestOp, (m, n) -> (n, m)} [id D] 0
│ └─ <Tensor3(float64, shape=(5, ?, ?))> [id E]
└─ Shape_i{2} [id F] 1
└─ Blockwise{MyTestOp, (m, n) -> (n, m)} [id D] 0
└─ ···
Without infer_shape in base Op:
MakeVector{dtype='int64'} [id A] 2
├─ 5 [id B]
├─ Shape_i{2} [id C] 1
│ └─ <Tensor3(float64, shape=(5, ?, ?))> [id D]
└─ Shape_i{1} [id E] 0
└─ <Tensor3(float64, shape=(5, ?, ?))> [id D]
For easy reference, my base infer_shape is as follows:
def infer_shape(self, fgraph, node, input_shapes):
if hasattr(self, "gufunc_signature"):
from pytensor.tensor.utils import _gufunc_to_out_shape
return _gufunc_to_out_shape(self.gufunc_signature, input_shapes)
else:
from pytensor.tensor.exceptions import ShapeError
raise ShapeError(f"Op {self} does not implement infer_shape")There was a problem hiding this comment.
Somehow the Blockwise shape is not being inferred correctly, the change by itself looks fine so we need to investigate more why does the Blockwise not disappear from the graph now. Does the Blockwise.infer_shape even get triggered now?
| @@ -202,6 +202,39 @@ def _parse_gufunc_signature( | |||
| ) | |||
|
|
|||
|
|
|||
| def _gufunc_to_out_shape( | |||
There was a problem hiding this comment.
This should try to prioritize input dimensions that are constant when there are multiple ones with the same letter, as it will generate a better shape graph.
Also fail explicitly if the shape can't be inferred from the signature alone
There was a problem hiding this comment.
Should I raise a ValueError when a single dimension in the gufunc_signature (say m) is assigned two values in the input shapes? I am raising a ValueError when the shape cannot be inferred from the signature. Here is my new function:
def _gufunc_to_out_shape(
signature: str, shapes: list[tuple[int, ...]]
) -> list[tuple[int, ...]]:
"""
Compute the shape of the output of an Op given its gufunc signature and the
shapes of its inputs.
Parameters
----------
signature : str
The gufunc signature of the Op.
eg: "(m,n),(n,p)->(m,p)".
shapes : list of tuple of int
The list of shapes of the inputs.
Returns
-------
out_shape : list of tuple of int
The list of shapes of the outputs.
Raises
------
ValueError
If the signature is invalid for the shapes of the inputs.
"""
parsed = _parse_gufunc_signature(signature)
out_shape = []
dim_to_size = dict()
for i, input_shape in enumerate(parsed[0]):
for j, dim in enumerate(input_shape):
if dim not in dim_to_size:
dim_to_size[dim] = shapes[i][j]
elif dim_to_size[dim] != shapes[i][j]:
raise ValueError(
f"Invalid signature {signature} for shapes {shapes}. "
f"Dimension {dim} is not consistent across inputs."
)
for i, output_shape in enumerate(parsed[1]):
temp_list = []
for j, dim in enumerate(output_shape):
if dim not in dim_to_size:
raise ValueError(
f"Invalid signature {signature} for shapes {shapes}. "
f"Dimension {dim} not in input dimensions."
)
else:
temp_list.append(dim_to_size[dim])
out_shape.append(tuple(temp_list))
return out_shapeThere was a problem hiding this comment.
shapes are symbolic you can't always compare them to know if they match, you can only do that for constants
There was a problem hiding this comment.
Okay, so if shapes list can contain symbolic values, can we first convert the symbolic variable to a tensor_variable using as_tensor_variables and then use .equals method to compare them? Specifically, in the function above, we make this change:
for i, input_shape in enumerate(parsed[0]):
for j, dim in enumerate(input_shape):
current_dim = as_tensor_variable(shapes[i][j])
if dim not in dim_to_size:
dim_to_size[dim] = current_dim
elif not dim_to_size[dim].equals(current_dim):
raise ValueError(
f"Invalid signature {signature} for shapes {shapes}. "
f"Dimension {dim} is not consistent across inputs."
)This is how the output looks like then: (The shapes input contains integers here)
print(_gufunc_to_out_shape("(m,m)->(m),(m,m)", [(2, 2)]))
[(TensorConstant(TensorType(int8, shape=()), data=array(2, dtype=int8)),), (TensorConstant(TensorType(int8, shape=()), data=array(2, dtype=int8)), TensorConstant(TensorType(int8, shape=()), data=array(2, dtype=int8)))]Can there be a better way to do this?
There was a problem hiding this comment.
No, if they are symbolic, unless they are constants, you can't know if they are equivalent, so raising an error is incorrect. User may have defined x = pt.vector("x"); y =pt.vector("y"); out = x + y. You will not know if x and y have identical shapes until the user compiles a function and provides values for them.
The logic for giving priority to constants is already in the Blockwise infer_shape, you should be able to just grab it and refactor it. We just want to simplify the graph returned, if there are two inputs with the same letters, and one of them has a constant shape. In that case we pick the constant one.
Checking if two inputs agree is not the critical thing here, although we can do that. You can only do that if isinstance(shapes[0][j], Constant) and isinstance(shapes[1][j], Constant) in which case you can then check shapes[0][j].data == shapes[1][j].data if you want to raise an informative error when they are inconsistent.
Also no reason to convert to as_tensor_variable, they should be ScalarVariables IIRC. Could be wrong here.
There was a problem hiding this comment.
I took the code from the infer_shape of Blockwise class, and used it to write the function. This gives priority to Constants. Is the output type list[tuple[Any, ...]] fine?
def _gufunc_to_out_shape(
signature: str, shapes
) -> list[tuple[Any, ...]]:
"""
Compute the shape of the output of an Op given its gufunc signature and the
shapes of its inputs.
Parameters
----------
signature : str
The gufunc signature of the Op.
eg: "(m,n),(n,p)->(m,p)".
shapes : list of tuple of Any
The list of shapes of the inputs.
Returns
-------
out_shape : list of tuple of Any
The list of shapes of the outputs.
Raises
------
ValueError
If the signature is invalid for the shapes of the inputs.
"""
input_sig, output_sig = _parse_gufunc_signature(signature)
dim_to_size : dict[str, Any] = {}
for input_shape, sig in zip(shapes, input_sig, strict = True):
for size, dim_name in zip(input_shape, sig, strict=True):
prev_size = dim_to_size.get(dim_name)
if prev_size is None:
dim_to_size[dim_name] = size
# Prefer constants
elif not isinstance(prev_size, Constant):
dim_to_size[dim_name] = size
elif prev_size.data != size:
raise ValueError(
f"Invalid signature {signature} for shapes {shapes}. "
f"Dimension {dim_name} is not consistent across inputs."
)
out_shapes = []
for output_shape in output_sig:
temp_list = []
for dim in output_shape:
if dim not in dim_to_size:
raise ValueError(
f"Invalid signature {signature} for shapes {shapes}. "
f"Dimension {dim} not in input dimensions."
)
else:
temp_list.append(dim_to_size[dim])
out_shapes.append((*temp_list,))
return out_shapes
…ut shape computation
…ons along with error handling
Aarsh-Wankar
force-pushed
the
infer_shape_1257
branch
from
3c8d955 to
870b900
Compare
…ntations in subclasses
…er_shape in Op class
Refactored
infer_shapemethod ofOpsto find output shapes usinggufunc_signatureusing a newly defined function_gufunc_to_out_shape.Description
Opshave a methodinfer_shapewhich helps to find the shapes of the outputs given the input shapes. AllOpshave their own implementations ofinfer_shape. However, manyOpshave agufunc_signaturestring which gives information about the input-output shape relations. In principle, this string is enough to find the output shapes given the input shapes. Thus, a function_gufunc_to_out_shapehas been added in this PR, which calculates the output shapes list given thegufunc_signatureand the input shapes list. This PR also replaces the Op specific implementations ofinfer_shapewith the_gufunc_to_out_shapeoutput.Related Issue
infer_shapeautomatically fromgufunc_signature#1257Checklist
Type of change
📚 Documentation preview 📚: https://pytensor--1294.org.readthedocs.build/en/1294/