Just as replaced in the new version of Opacus, here we switch to using einsum from https://optimized-einsum.readthedocs.io/en/stable/ that can significantly reduce the overall execution time of einsum-like expressions by optimizing the expression’s contraction order and dispatching many operations to canonical BLAS, cuBLAS, or other specialized routines (quoted from the website).
Just as replaced in the new version of Opacus, here we switch to using einsum from https://optimized-einsum.readthedocs.io/en/stable/ that can significantly reduce the overall execution time of einsum-like expressions by optimizing the expression’s contraction order and dispatching many operations to canonical BLAS, cuBLAS, or other specialized routines (quoted from the website).