Function func

imperative/python/test/unit/xla/test_xla_trace.py:22–24

(x, y)

@pytest.mark.skipif(not is_cuda_available(), reason="only support cuda now")
def test_partial_trace_inplace():
    def func(x, y):
        x += 1
        y += 1

    xla_func = partial_trace(func, backend="xla")
    xla_func(tensor(1), tensor(2))