einsumfunc.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536
  1. from __future__ import annotations
  2. from typing import Any
  3. import numpy as np
  4. AR_LIKE_b = [True, True, True]
  5. AR_LIKE_u = [np.uint32(1), np.uint32(2), np.uint32(3)]
  6. AR_LIKE_i = [1, 2, 3]
  7. AR_LIKE_f = [1.0, 2.0, 3.0]
  8. AR_LIKE_c = [1j, 2j, 3j]
  9. AR_LIKE_U = ["1", "2", "3"]
  10. OUT_f: np.ndarray[Any, np.dtype[np.float64]] = np.empty(3, dtype=np.float64)
  11. OUT_c: np.ndarray[Any, np.dtype[np.complex128]] = np.empty(3, dtype=np.complex128)
  12. np.einsum("i,i->i", AR_LIKE_b, AR_LIKE_b)
  13. np.einsum("i,i->i", AR_LIKE_u, AR_LIKE_u)
  14. np.einsum("i,i->i", AR_LIKE_i, AR_LIKE_i)
  15. np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f)
  16. np.einsum("i,i->i", AR_LIKE_c, AR_LIKE_c)
  17. np.einsum("i,i->i", AR_LIKE_b, AR_LIKE_i)
  18. np.einsum("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c)
  19. np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f, dtype="c16")
  20. np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=bool, casting="unsafe")
  21. np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f, out=OUT_c)
  22. np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=int, casting="unsafe", out=OUT_f)
  23. np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_b)
  24. np.einsum_path("i,i->i", AR_LIKE_u, AR_LIKE_u)
  25. np.einsum_path("i,i->i", AR_LIKE_i, AR_LIKE_i)
  26. np.einsum_path("i,i->i", AR_LIKE_f, AR_LIKE_f)
  27. np.einsum_path("i,i->i", AR_LIKE_c, AR_LIKE_c)
  28. np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_i)
  29. np.einsum_path("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c)