Description

numpy compatible einsum

Changes

• add einsum
• add einsum doc
• add path optimization
• add benchmark
• add forward and backward test

Comments

• support both explicit and implicit mode
• support multiple operands
• support broadcast
• support 0-dim NDArray and 0-size NDArray
• cannot work until multiple tempspace is enabled
  • einsum make use of multiple temspace, but set MXNET_GPU_TEMP_COPY to 4 triggers a bug
  • simply to set MXNET_GPU_TEMP_COPY to 4 (without the presence of einsum) triggers the bug. So this might be a bug of mxnet itself.
  • solution 1: fix the bug (may be difficult)
  • solution 2: implement einsum in another way and sidestep this feature.

TODO

• translate einsum_path into c, so that optimization can work in symbol mode
• take advantage of tensordot when it is merged
• support "greedy" optimization
• support "optimal" optimization

Thank @reminisce and @haojin2 for guidance and review.