-
Notifications
You must be signed in to change notification settings - Fork 136
Closed
Labels
Description
Description
Revealed in #1226
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.tensor.slinalg import SolveTriangular
rng = np.random.default_rng(1)
a_test = rng.normal(size=(2, 2))
a_test[1, 0] = 0
b_test = rng.normal(size=(2, 2))
op = SolveTriangular(trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2, overwrite_b=True)
a_pt = pt.matrix("a", shape=a_test.shape)
b_pt = pt.matrix("b", shape=b_test.shape)
out = op(a_pt, b_pt)
py_fn = pytensor.function([a_pt, b_pt], out, accept_inplace=True)
numba_fn = pytensor.function([a_pt, b_pt], out, accept_inplace=True, mode="NUMBA")
(
py_fn(a_test.copy(), b_test.copy()),
numba_fn(a_test.copy(), b_test.copy()),
)
# (array([[ 1.64016796, 2.35184226],
# [ 0.41204025, -0.44593092]]),
# array([[ 3.4341482 , -0.34253317],
# [-0.49356511, -0.44593092]]))