Describe the bug
Residuals computed for each equation in SystemEquation are expected to be concatenated along the last dimension. The current implementation uses torch.hstack, which concatenates along dim=1 for tensors with ndim ≥ 2. While this works for 2D tensors, it produces incorrect behavior for tensors with more than two dimensions, as it forces all dimensions except dim=1 to match. As a result, shape inconsistencies arise when operating on higher-dimensional tensors.
Expected behavior
Residuals should be concatenated along the last dimension, ensuring compatibility with tensors of arbitrary dimensionality.
Additional context
The issue can be resolved by replacing the current stacking operation (hstack) with a concatenation explicitly performed along the last dimension (cat(..., dim=-1)). This ensures consistent behavior across tensors of arbitrary dimension and prevents shape mismatches when working with higher-dimensional inputs.
Describe the bug
Residuals computed for each equation in
SystemEquationare expected to be concatenated along the last dimension. The current implementation usestorch.hstack, which concatenates alongdim=1for tensors withndim ≥ 2. While this works for 2D tensors, it produces incorrect behavior for tensors with more than two dimensions, as it forces all dimensions exceptdim=1to match. As a result, shape inconsistencies arise when operating on higher-dimensional tensors.Expected behavior
Residuals should be concatenated along the last dimension, ensuring compatibility with tensors of arbitrary dimensionality.
Additional context
The issue can be resolved by replacing the current stacking operation (
hstack) with a concatenation explicitly performed along the last dimension (cat(..., dim=-1)). This ensures consistent behavior across tensors of arbitrary dimension and prevents shape mismatches when working with higher-dimensional inputs.