You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
For long term maintenance simplicity, consider replacing the custom cuda kernels redrock.zscan.batch_dot_product_3d3d and batch_dot_product_3d2d with einsum magic as suggested by @dmargala :
For example, batched A.T.dot(A) and A.T.dot(b) would be:
cp.einsum("...ji,...jk", A, A)
cp.einsum("...ji,...j", A, b)
Those aren't a drop-in replacement for the call signature of batch_dot_product_3d3d, but I think we are using it for that A.T.dot(A) purpose. Profile test it against current implementation and also check for correctness.
Also consider moving functions like this into redrock.utils or a separate redrock.linalg or similar module instead of zscan.
The text was updated successfully, but these errors were encountered:
For long term maintenance simplicity, consider replacing the custom cuda kernels
redrock.zscan.batch_dot_product_3d3d
andbatch_dot_product_3d2d
witheinsum
magic as suggested by @dmargala :Those aren't a drop-in replacement for the call signature of
batch_dot_product_3d3d
, but I think we are using it for thatA.T.dot(A)
purpose. Profile test it against current implementation and also check for correctness.Also consider moving functions like this into
redrock.utils
or a separateredrock.linalg
or similar module instead of zscan.The text was updated successfully, but these errors were encountered: