pt.linalg.kron destroys static shape information

[jessegrabowski]

Description

A = pt.dmatrix("a", shape=(4,2))
B = pt.dmatrix('b', shape=(3,3))
z = pt.linalg.kron(A, B)
z.type.shape #(None, None)

I expect z.type.shape to be (12, 6), since all input shapes are known

[Utkarshkarki]

The issue can be addressed by adding conditional shape inference to pt.linalg.kron. When both inputs are rank-2 tensors (matrices) and their shapes are fully known, the output shape is mathematically deterministic and can be inferred as (m·p, n·q). In all other cases—such as dynamic shapes or higher-rank tensors—the operation should safely fall back to an unknown shape. This preserves correctness while allowing PyTensor to retain useful static shape information whenever it is available.

[Utkarshkarki]

please tell if i should go ahead with this or not

[jessegrabowski]

Hey @Utkarshkarki , the issue is here. out_shape is using strictly symbolic values a.shape, b.shape. Instead, we should use the known static shapes if they are known, otherwise the symbolic shapes. Something like this.

[ricardoV94]

But that's an OpFromGraph, you don't have access to the make_node. The issue here is that reshape can't figure out the static shape. You could add a manual specify_shape from the static shapes I suppose

[jessegrabowski]

I didn't literally mean there should be a make_node. I meant that the shapes passed to out_shape can have the same conditionally known logic as is in the linked make node. This reshape is the step that destroys the static shape information, because it only ever gets static shapes. If we pass static shapes, they will be correctly set.

[ricardoV94]

Do note: Every time you fix a static shape symbolically (e.g., by inlining x.type.shape) you reduce the chances of doing an "unsafe" clone_replace with different shaped inputs, because you broke the symbolic source of that shape.

It's sort of like an .eval() for shapes

[jessegrabowski]

You mean increase?

In that case a deeper fix would be to make dummy inputs, do the computation on the core case, Blockwise the core OpFromGraph, then pass in the actual user inputs? I think all the static shape propagation would just work in that case.

[ricardoV94]

I mean reduce. Right now it's possible to do:

x = pt.tensor(shape=(5,))
y = pt.zeros(x.shape)  
# y = pt.zeros(x.type.shape)
clone_replace(y, {x: pt.tensor(shape=(6,))})

If you use x.type.shape, you lose that option.

The ask in this PR is fine, but just stating that depending on what you want to do, you can get trapped by eagerly "constant-folding" symbolic shapes, and it's easy to want to do that when working with Reshape

[jessegrabowski]

Well let's try to ask for a principled solution in this PR. If we have something like:

def my_ofg(a, b):
   a_core = pt.dmatrix(shape=a.shape[-2:])
   b_core = pt.dmatrix(shape=b.shape[-2:])
   result_core = fn(a_core, b_core)
   op = OFG(inputs=[a_core, b_core], outputs=[result_core])

   return Blockwise(op)(a, b)

Assume OFG has a gufunc_signature. Does this do what we want?

[ricardoV94]

I imagine that has exactly the same static shape loss?

[jessegrabowski]

does the linked make_node ruin clone_replace?

[ricardoV94]

You linked to two cases. KroneckerProduct is a subclass of OpFromGraph, and BlockDiagonal which is a regular Op.

With regular Ops you have complete control of the output types (you must define them).

OpFromGraph assumes your output types will be the ones from the output variables. If you know better, you can add a specify_shape, so the output type will be more informative.

It may have those lock-in consequences though.

Otherwise you could override OpFromGraph.make_node but I advise you against that.