Skip to content

Instantly share code, notes, and snippets.

@deanmark
deanmark / tensordot_pytorch.py
Created June 12, 2018 10:59
tensordot in pytorch
def tensordot_pytorch(a, b, axes=2):
# code adapted from numpy
try:
iter(axes)
except Exception:
axes_a = list(range(-axes, 0))
axes_b = list(range(0, axes))
else:
axes_a, axes_b = axes
try: