BUG: `tensordot` is broken about pytensor

[ricardoV94]

the type.shape are not propagated properly

This is a nice have but not a must. PyTensor is only requires to propagate broadcastable info at write time.

In any case, the loss likely comes from the Reshape Op.

The complicated tensordot should probably be a rewrite, as I'm sure JAX has a nice support for it without reshape tricks. Numba I don't know.

Neither of these preclude just patching the current bug.

[ricardoV94]

the type.shape are not propagated properly

This is a nice have but not a must. PyTensor is only required to propagate broadcastable info at write time.

In any case, the loss likely comes from the Reshape Op.

The complicated tensordot should probably be a rewrite, as I'm sure JAX has a nice support for it without reshape tricks. Numba I don't know.

Neither of these preclude just patching the current bug.

[lucianopaz]

I found that I can mostly copy over the implementation from numpy. The only stuff that I had to do differently is handling broadcastable, x.type.shape and x.shape properly to ensure everything runs smoothly.

The problem is that numpy has no implementation for batched, and to be honest, it seems to me like keeping batched versions of dot and tensordot is completely pointless. @ricardoV94, what do you think?

[ricardoV94]

I don't recall exactly what batched was doing. It uses BatchedDot to create a Batched TensorDot or something along those lines?

[lucianopaz]

It treats the first dimension of the tensors as batches. Mostly like a crude Blockwise

[ricardoV94]

It treats the first dimension of the tensors as batches. Mostly like a crude Blockwise

If that is not being used anywhere in the codebase we can remove with perhaps an informative error message if anyone tries to call with the old signature. I think we can handle that better by trying to optimize Blockwise dots at compile time instead.

[ricardoV94]

Which dimension was assumed to be batched, the leading one? I see you mentioned it.

For backwards-compat you could also just use pt.vectorize after doing the core tensordot (may or not need to change the input axis), and still start deprecating the kwarg (in this case legacy code wouldn't be broken immediately, just probably run slower)

Not sure it's worth it!