tensor-sensor

1
import tsensor
2
import numpy as np
3

4
def f():
5
    # Currently can't handle double assign
6
    a = b = np.ones(1) @ np.ones(2)
7

8
def A():
9
    with tsensor.clarify():
10
        f()
11

12

13
def test_nested():
14
    msg = ""
15
    try:
16
        A()
17
    except BaseException as e:
18
        msg = e.args[0]
19

20
    expected = "matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 2 is different from 1)"
21
    assert msg==expected
22