def test_type_tracking():
pip = scalar_pipeline.with_steps(
steps.step_parse,
steps.step_infer,
steps.step_specialize,
steps.step_simplify_types,
LocalPassOptimizer(opt_ok1, opt_ok2, opt_err1),
steps.step_validate,
)
def fn_ok1(x, y):
return x + y
pip(input=fn_ok1, argspec=(to_abstract_test(i64), to_abstract_test(i64)))
def fn_ok2(x):
return -x
pip(input=fn_ok2, argspec=(to_abstract_test(i64),))
def fn_err1(x, y):
return x - y
with pytest.raises(ValidationError):
pip(
input=fn_err1,
argspec=(to_abstract_test(i64), to_abstract_test(i64)),
)
def test_incorporate_call_through_switch():
def before_helper(x):
if x < 0:
return scalar_mul
else:
return scalar_add
def before(x, y, z):
return before_helper(x)(y, z)
def after(x, y, z):
def after_helper(x, y, z):
def tb(y, z):
return scalar_mul(y, z)
def fb(y, z):
return scalar_add(y, z)
return switch(x < 0, tb, fb)(y, z)
return after_helper(x, y, z)
_check_opt(
before,
after,
lib.elim_identity,
lib.incorporate_call,
lib.incorporate_call_through_switch,
argspec=[
to_abstract_test(f64),
to_abstract_test(f64),
to_abstract_test(f64),
],
)
def test_incorporate_env_getitem_through_switch():
def before(x, y):
key = embed(x)
def f1(x, y):
return env_setitem(newenv, key, x * y)
def f2(x, y):
return env_setitem(newenv, key, x + y)
return env_getitem(switch(x < 0, f1, f2)(x, y), key, 0)
def after(x, y):
def f1(x, y):
return x * y
def f2(x, y):
return x + y
return switch(x < 0, f1, f2)(x, y)
_check_opt(
before,
after,
lib.incorporate_env_getitem_through_switch,
lib.cancel_env_set_get,
argspec=[to_abstract_test(f64),
to_abstract_test(f64)],
)
def test_incorporate_getitem_through_switch():
def before(x, y):
def f1(x, y):
return x, y
def f2(x, y):
return y, x
return switch(x < 0, f1, f2)(x, y)[0]
def after(x, y):
def f1(x, y):
return x
def f2(x, y):
return y
return switch(x < 0, f1, f2)(x, y)
_check_opt(
before,
after,
lib.incorporate_getitem_through_switch,
argspec=[to_abstract_test(f64),
to_abstract_test(f64)],
)
def test_incorporate_call():
def b_help(q):
def subf(z):
return q * z
return subf
def before(x, y):
return b_help(x)(y)
def after(x, y):
def a_help(q, y):
def subf(z):
return q * z
return subf(y)
return a_help(x, y)
_check_opt(
before,
after,
lib.incorporate_call,
argspec=[to_abstract_test(f64),
to_abstract_test(f64)],
)
def _grad_test(
fn,
obj,
args,
sens_type=f64,
pipeline=grad_pipeline,
rel_error=1e-3,
argspec=None,
):
pipeline = pipeline.insert_after(steps.step_parse, grad_wrap)
if argspec is None:
argspec = tuple(
from_value(arg, broaden=True) for arg in clean_args(args))
else:
argspec = tuple(to_abstract_test(x) for x in argspec)
sens_type = to_abstract_test(sens_type)
if isinstance(obj, FunctionType):
res = pipeline(input=obj, argspec=[*argspec, sens_type])
else:
pip = pipeline.without_step(steps.step_parse)
res = pip(graph=obj, argspec=[*argspec, sens_type])
gtest = GradTester(
fn=fn,
gfn=res["output"],
args=args,
argnames=[f"in{i}" for i in range(len(args))],
outnames=None,
rel_error=rel_error,
)
gtest.assert_match()
def test_validate_abstract_2():
bad_array = AbstractArray(to_abstract_test(f64), {
SHAPE: (1, 2),
TYPE: PyTorchTensor
})
with pytest.raises(ValidationError):
validate_abstract(bad_array, uses={})
def test_env_get_set():
def before(x, y):
a = 5678
e = env_setitem(newenv, embed(x), y)
e = env_setitem(e, embed(a), a)
return env_getitem(e, embed(x), 1234)
def after(x, y):
return y
_check_opt(
before,
after,
lib.cancel_env_set_get,
argspec=[to_abstract_test(f64),
to_abstract_test(f64)],
)
def test_type_tracking_2():
pip = scalar_pipeline.with_steps(
steps.step_parse,
steps.step_infer,
steps.step_specialize,
steps.step_simplify_types,
LocalPassOptimizer(opt_ok1, opt_ok2, opt_err1),
steps.step_validate,
)
def fn_err3(x, y):
return x - y + x
with pytest.raises(InferenceError):
pip(
input=fn_err3,
argspec=(to_abstract_test(i64), to_abstract_test(i64)),
)
def test_incorporate_getitem_2():
def before(x, y):
def b_help(x, y):
return x
return b_help(x, y)[0]
def after(x, y):
def a_help(x, y):
return x[0]
return a_help(x, y)
_check_opt(
before,
after,
lib.incorporate_getitem,
argspec=[to_abstract_test((f64, f64)),
to_abstract_test(f64)],
)
def test_getitem_newenv():
def before(x):
return env_getitem(newenv, embed(x), 1234)
def after(x):
return 1234
_check_opt(before,
after,
lib.getitem_newenv,
argspec=[to_abstract_test(f64)])
def test_env():
def f(x, y):
e1 = env_setitem(newenv, embed(x), 100)
e2 = env_setitem(newenv, embed(x), 10)
e2 = env_setitem(e2, embed(y), 20)
e3 = env_add(e1, e2)
a = env_getitem(e3, embed(x), 0)
b = env_getitem(e3, embed(y), 0)
c = env_getitem(e3, embed(a), 0)
return (a, b, c)
res = scalar_debug_pipeline(input=f,
argspec=(to_abstract_test(i64),
to_abstract_test(i64)))["output"](3,
4)
assert res == (110, 20, 0)
def test_incorporate_env_getitem_2():
def before(x, y):
def b_help(x, y):
return x
return env_getitem(b_help(x, y), embed(y), 0)
def after(x, y):
def a_help(x, y):
return env_getitem(x, embed(y), 0)
return a_help(x, y)
_check_opt(
before,
after,
lib.incorporate_env_getitem,
lib.cancel_env_set_get,
argspec=[to_abstract_test(newenv),
to_abstract_test(f64)],
)
def test_env_get_add():
def before(x, y):
e1 = env_setitem(newenv, embed(x), x)
e1 = env_setitem(e1, embed(y), y)
e2 = env_setitem(newenv, embed(y), y)
e2 = env_setitem(e2, embed(x), x)
return env_getitem(env_add(e1, e2), embed(x), 0)
def after(x, y):
return gadd(x, x)
_check_opt(
before,
after,
lib.getitem_env_add,
lib.cancel_env_set_get,
argspec=[to_abstract_test(i64),
to_abstract_test(i64)],
argspec_after=False,
)
def test_repr():
s1 = to_abstract_test(1)
assert repr(s1) == "AbstractScalar(Int[64] = 1)"
s2 = to_abstract_test(f32)
assert repr(s2) == "AbstractScalar(Float[32])"
t1 = to_abstract_test((1, f32))
assert repr(t1) == f"AbstractTuple((Int[64] = 1, Float[32]))"
a1 = to_abstract_test(af32_of(4, 5))
assert repr(a1) == f"AbstractArray(Float[32] x 4 x 5)"
p1 = to_abstract_test(Point(1, f32))
assert repr(
p1) == f"AbstractClass(Point(x :: Int[64] = 1, y :: Float[32]))"
j1 = AbstractJTagged(to_abstract_test(1))
assert repr(j1) == f"AbstractJTagged(J(Int[64] = 1))"
h1 = AbstractHandle(to_abstract_test(1))
assert repr(h1) == f"AbstractHandle(H(Int[64] = 1))"
kw1 = AbstractKeywordArgument("bucket", to_abstract_test(1))
assert repr(kw1) == f"AbstractKeywordArgument(KW(bucket :: Int[64] = 1))"
ty1 = Ty(f32)
assert repr(ty1) == "AbstractType(Ty(AbstractScalar(Float[32])))"
e1 = AbstractError(DEAD)
assert repr(e1) == "AbstractError(E(DEAD))"
f1 = AbstractFunction(P.scalar_mul)
assert repr(f1) == "AbstractFunction(scalar_mul)"
fa = AbstractFunction(value=ANYTHING)
assert repr(fa) == "AbstractFunction(ANYTHING)"
tu1 = AbstractTaggedUnion([[13, s2], [4, to_abstract_test(i16)]])
assert repr(tu1) == "AbstractTaggedUnion(U(4 :: Int[16], 13 :: Float[32]))"
bot = AbstractBottom()
assert repr(bot) == "AbstractBottom(⊥)"
def test_incorporate_env_getitem():
def before(x, y):
key = embed(x)
def b_help(x, y):
return env_setitem(newenv, key, x * y)
return env_getitem(b_help(x, y), key, 0)
def after(x, y):
def a_help(x, y):
return x * y
return a_help(x, y)
_check_opt(
before,
after,
lib.incorporate_env_getitem,
lib.cancel_env_set_get,
argspec=[to_abstract_test(f64),
to_abstract_test(f64)],
)
def test_build_value():
assert build_value(S(1)) == 1
with pytest.raises(ValueError):
build_value(S(t=ty.Int[64]))
assert build_value(S(t=ty.Int[64]), default=ANYTHING) is ANYTHING
assert build_value(T([S(1), S(2)])) == (1, 2)
loop = InferenceLoop(errtype=Exception)
p = loop.create_pending(resolve=(lambda: None), priority=(lambda: None))
with pytest.raises(ValueError):
assert build_value(S(p, t=ty.Int[64])) is p
assert build_value(S(p, t=ty.Int[64]), default=ANYTHING) is ANYTHING
p.set_result(1234)
assert build_value(S(p, t=ty.Int[64])) == 1234
pt = Point(1, 2)
assert build_value(to_abstract_test(pt)) == pt
def test_to_canonical():
def _convert(data, typ):
return to_canonical(data, to_abstract_test(typ))
# Leaves
assert _convert(True, Bool) is True
assert _convert(False, Bool) is False
assert _convert(10, i64) == 10
assert _convert(1.5, f64) == 1.5
assert _convert([], []) == ()
with pytest.raises(TypeError):
_convert([], [f64])
with pytest.raises(TypeError):
_convert([1, 2], [])
# Class -> Tuple conversion
pt = Point(1, 2)
pt3 = Point3D(1, 2, 3)
assert list(_convert(pt, Point(i64, i64))) == [1, 2]
with pytest.raises(TypeError):
_convert((1, 2), Point(i64, i64))
assert list(_convert((pt, pt), (Point(i64, i64), Point(i64, i64)))) == [
(1, 2),
(1, 2),
]
li = _convert([1], [i64])
assert (isinstance(li, tuple) and li[0] == 1
and isinstance(li[1], TaggedValue) and li[1].value == ())
# Arrays
fmat = np.ones((5, 8))
imat = np.ones((5, 8), dtype="int32")
assert _convert(fmat, af64_of(5, 8)) is fmat
assert _convert(imat, ai32_of(5, 8)) is imat
with pytest.raises(TypeError):
_convert(imat, ai64_of(5, 8))
with pytest.raises(TypeError):
_convert(imat, ai32_of(4, 8))
# Misc errors
with pytest.raises(TypeError):
_convert(10, f64)
with pytest.raises(TypeError):
_convert(1.5, i64)
with pytest.raises(TypeError):
_convert(10, (i64, i64))
with pytest.raises(TypeError):
_convert((1, ), (i64, i64))
with pytest.raises(TypeError):
_convert((1, 2, 3), (i64, i64))
with pytest.raises(TypeError):
_convert((1, 2, 3), [i64])
with pytest.raises(TypeError):
_convert(pt3, Point(i64, i64))
with pytest.raises(TypeError):
_convert(10, ai64_of())
with pytest.raises(TypeError):
_convert(10, ai64_of())
with pytest.raises(TypeError):
_convert(1, Bool)
with pytest.raises(TypeError):
_convert(1, D(x=i64))
with pytest.raises(TypeError):
_convert({"x": 2.0}, D(x=i64))
with pytest.raises(TypeError):
_convert({"x": 2.0, "y": 1}, D(x=i64))
with pytest.raises(TypeError):
_convert({"y": 2.0}, D(x=i64))
with pytest.raises(TypeError):
_convert("x", 1.0)
with pytest.raises(TypeError):
_convert(1.0, to_abstract_test("x"))
with pytest.raises(TypeError):
_convert(1.0, to_abstract_test(HandleInstance(1.0)))
v = to_device(22, None)
with pytest.raises(TypeError):
_convert(v, f64)
from myia.lib import InferenceError
from myia.operations import dtype, scalar_cast
from myia.testing.common import MA, Ty, af32_of, f32, i64, to_abstract_test
from myia.testing.multitest import infer, mt, run
@mt(
infer(i64, result=InferenceError),
infer(af32_of(4, 5), result=Ty(to_abstract_test(f32))),
)
def test_dtype(arr):
return dtype(arr)
@mt(infer(af32_of(4, 5), i64, result=f32), run(MA(2, 3), 7, result=7.0))
def test_cast_to_dtype(arr, x):
return scalar_cast(x, dtype(arr))
def run(pip, fn, types):
res = pip(input=fn, argspec=[to_abstract_test(t) for t in types])
return res["graph"]
def test_validate_abstract():
scalar_i64 = to_abstract_test(i64)
fn = AbstractFunction(
TypedPrimitive(P.scalar_add, (scalar_i64, scalar_i64), scalar_i64), )
with pytest.raises(ValidationError):
validate_abstract(fn, uses={})
def _convert(data, typ):
return from_canonical(data, to_abstract_test(typ))
def test_numpy_scalar_to_abstract():
s1 = AbstractScalar({VALUE: 2, TYPE: i16})
assert to_abstract_test(np.int16(2)) == s1
s2 = AbstractScalar({VALUE: 1.5, TYPE: f32})
assert to_abstract_test(np.float32(1.5)) == s2
def _fwd_and_bwd(
self,
fn,
args,
broad_specs=None,
pipeline=standard_pipeline,
backend=False,
numpy_compat=True,
atol=1e-8,
rtol=1e-5,
grad_atol=1e-6,
grad_rtol=1e-5,
):
if backend:
backend_name = backend[0]
backend_options = backend[1]
pipeline = pipeline.configure({
"backend.name": backend_name,
"backend.options": backend_options
})
def mksens(x):
return AbstractArray(
AbstractScalar({
TYPE: pytorch_dtype_to_type(x.dtype),
VALUE: ANYTHING
}),
{
SHAPE: tuple(x.shape),
TYPE: PyTorchTensor
},
)
ref_result = fn(*map(copy, args))
argspec = make_argspec(args, broad_specs)
res = pipeline(input=fn, argspec=argspec)
myia_fn = res["output"]
myia_result = myia_fn(*map(copy, args))
assert eqtest(ref_result, myia_result, atol=atol, rtol=rtol)
if isinstance(myia_result, tuple):
sens_type = AbstractTuple([mksens(res) for res in myia_result])
sens = tuple(_make_sens(res) for res in myia_result)
else:
sens_type = mksens(myia_result)
sens = _make_sens(myia_result)
pytorch_grads = pt_fn_grads(fn, *args)
gpipeline = pipeline.insert_after(steps.step_parse, grad_wrap)
sens_type = to_abstract_test(sens_type)
assert isinstance(fn, FunctionType)
res = gpipeline(input=fn, argspec=[*argspec, sens_type])
myia_grads = res["output"](*args, sens)
assert eqtest(pytorch_grads, myia_grads, rtol=grad_rtol, atol=grad_atol)
if numpy_compat:
args_torch = args
args = ()
for _ in args_torch:
if isinstance(_, torch.Tensor):
args += (_.detach().numpy(), )
else:
args += (_, )
if backend:
backend_name = backend[0]
backend_options = backend[1]
pipeline = pipeline.configure({
"backend.name": backend_name,
"backend.options": backend_options,
})
def mksens(x):
return AbstractArray(
AbstractScalar({
TYPE: np_dtype_to_type(x.dtype.name),
VALUE: ANYTHING
}),
{
SHAPE: tuple(x.shape),
TYPE: NDArray
},
)
argspec = make_argspec(args, broad_specs)
res = pipeline(input=fn, argspec=argspec)
myia_fn = res["output"]
myia_result = myia_fn(*map(copy, args))
if isinstance(myia_result, tuple):
sens_type = AbstractTuple([mksens(res) for res in myia_result])
sens = tuple(_make_sens_numpy(res) for res in myia_result)
else:
sens_type = mksens(myia_result)
sens = _make_sens_numpy(myia_result)
gpipeline = pipeline.insert_after(steps.step_parse, grad_wrap)
sens_type = to_abstract_test(sens_type)
assert isinstance(fn, FunctionType)
res = gpipeline(input=fn, argspec=[*argspec, sens_type])
myia_grads = res["output"](*args, sens)
