Fix a bug of forward-mode AD when multi-output is needed

[Jerry-Jzy]

I find the bug from this case: https://github.com/lululxvi/deepxde/blob/master/examples/operator/stokes_aligned_pideeponet.py

when setting the num_output=3 and multi_output_strategy="independent", the output shape will be (batch size, # of coordinates, 3).
Then, the slice here is wrong, what we really want is slicing in the last dim.

[lululxvi]

Instead of doing reshape here, another solution is modifying dde.gradients to support the output shape of (batch 1, batch 2, dim). This might be a better solution.

[Jerry-Jzy]

Instead of doing reshape here, another solution is modifying dde.gradients to support the output shape of (batch 1, batch 2, dim). This might be a better solution.

I think reshape in PDEOperatorCartesianProd is also a good solution. Since the problem is originally from the output structure difference between PI-DeepONet and PINN. After unifying the output structure, dde.gradients is a general-purpose module for gradent computation

[lululxvi]

Can you test another problem?
change this line

deepxde/examples/operator/advection_aligned_pideeponet.py

Line 33 in b313d96

return dy_t + dy_x

to

return dy_t + dy_x + x[:, 0:1]

The reverse autodiff should run, but I am not sure about the current forward code.

[Jerry-Jzy]

Can you test another problem? change this line

deepxde/examples/operator/advection_aligned_pideeponet.py

Line 33 in b313d96

return dy_t + dy_x

to

return dy_t + dy_x + x[:, 0:1]

The reverse autodiff should run, but I am not sure about the current forward code.

the forward-mode is not suitable for this case. The reason is not reshape. Since the forward-code is loop-free. so the shape gradient is same as the output shape. The reverse mode is doable because loop is needed then the shape of gradient is same as input shape. That's the exact difference between JVP and VJP. If we want this case working. we need to repeat the x num_func times.

[lululxvi]

That is why I consider to modify dde.gradients to handle and keep 3D shape; in this way, we can use broadcast, so no repeating x is needed.

[Jerry-Jzy]

That is why I consider to modify dde.gradients to keep 3D shape; in this way, we can use broadcast, so no repeating x is needed.

although 3D shape is kept, Still cannot use broadcast. since the (num_func, num_points) can not broadcast with (num_points, 1).

[lululxvi]

(num_func, num_points, 1) can broadcast with (num_points, 1)

[lululxvi]

return dy_t + dy_x + x[:, 0:1]

The reverse autodiff should run, but I am not sure about the current forward code.

The key question now is that if the current reshape approach can solve this problem without asking users to repeat x. If not, then we need to find out another way.

[lululxvi]

Is current code working for single output? If so, let us make the code work for single output, and just raise an error for multi-outputs. Then we will return to multi-output in another PR.

[Jerry-Jzy]

(num_func, num_points, 1) can broadcast with (num_points, 1)

If we uniformly reshape to 3D rather than 2D maybe can solve the problem of this case

[lululxvi]

Is current code working for single output? If so, let us make the code work for single output, and just raise an error for multi-outputs. Then we will return to multi-output in another PR.

How about single output case?

[Jerry-Jzy]

Is current code working for single output? If so, let us make the code work for single output, and just raise an error for multi-outputs. Then we will return to multi-output in another PR.

How about single output case?

single output works fine

[lululxvi]

Then clean the code only for single output. We merge the single output code first. Or the current code in the main branch is good?

[Jerry-Jzy]

Then clean the code only for single output. We merge the single output code first. Or the current code in the main branch is good?

I have rolled back the code