def assert_eq(a, b):
    if hasattr(a, "compute"):
        a = a.compute(scheduler="sync")
    if hasattr(b, "compute"):
        b = b.compute(scheduler="sync")

    assert a == b