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)