MIONet for pytorch backend

[voloddia]

Hi Professor Lu,
Currently I see there is no MIONet class for pytorch backend.
By looking at the code for tensorflow.compat.v1, I guess I need to replace

if self.output_merger is None:
    y = torch.einsum("ip,jp->ij", y_func, y_loc)

with

if self.output_merger is None:
    y = torch.mul(y_func, y_loc)
    y = torch.sum(y, axis=1, keepdim=True)

But I'm not sure if other things need to be modified when self.output_merger is not None.
Could you please advise?

[lululxvi]

Looks OK.

[voloddia]

Thank you, Professor Lu.
When self.output_merger is not None, I replace

# MIONetCartesianProd
else:
    y_func = y_func[:, None, :]
    y_loc = y_loc[None, :]
    if self.output_merge_operation == "mul":
        y = torch.mul(y_func, y_loc)
    elif self.output_merge_operation == "add":
        y = y_func + y_loc
    elif self.output_merge_operation == "cat":
        y_func = y_func.repeat(1, y_loc.shape[1], 1)
        y_loc = y_loc.repeat(y_func.shape[0], 1, 1)
        y = torch.cat((y_func, y_loc), dim=2)
    shape0 = y.shape[0]
    shape1 = y.shape[1]
    y = y.reshape(shape0 * shape1, -1)
    y = self.output_merger(y)
    y = y.reshape(shape0, shape1)

with

# MIONet
else:
    if self.output_merge_operation == "mul":
        y = torch.mul(y_func, y_loc)
    elif self.output_merge_operation == "add":
        y = y_func + y_loc
    elif self.output_merge_operation == "cat":    
        y = torch.cat((y_func, y_loc), dim=1)    
    y = self.output_merger(y)

Is it correct?

[lululxvi]

Feel free to submit a PR to add MIONet.