Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[bug] Fix MatrixFreeCG so it can handle multiple input sizes. (#8070)
### Brief Summary Previously, `MatrixFreeCG` inappropriately assume that the intermediate vectors used in CG solver are all two-dimensional `ti.field`. This is incorrect because the dimension of `Ap` `p` and `r` should be consistent with input `b` and `x`. ### Walkthrough Instead of simply put `vector_fields_builder.dense(ti.ij, size).place(p, r, Ap)`, now it's based on the length of `size`. ```python def MatrixFreeCG(A, b, x, tol=1e-6, maxiter=5000, quiet=True): ... vector_fields_builder = ti.FieldsBuilder() p = ti.field(dtype=solver_dtype) r = ti.field(dtype=solver_dtype) Ap = ti.field(dtype=solver_dtype) if len(size) == 1: # Determine the `axes` argument based on the length of `size` axes = ti.i elif len(size) == 2: axes = ti.ij elif len(size) == 3: axes = ti.ijk else: raise TaichiRuntimeError(f"MatrixFreeCG currently cannot support {len(size)}-D inputs.") vector_fields_builder.dense(axes, size).place(p, r, Ap) ```
- Loading branch information