def test_remove_ids(): assert remove_ids(MatMul(A, Identity(m), B, evaluate=False)) == \ MatMul(A, B, evaluate=False) assert null_safe(remove_ids)(MatMul(Identity(n), evaluate=False)) == \ MatMul(Identity(n), evaluate=False)