wehub-resource-sync
1331 lines
66 KiB
Python
1331 lines
66 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E501, F841
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.ir.base import assert_structural_equal
|
|
from tvm.script.parser import ir as I
|
|
from tvm.script.parser import relax as R
|
|
from tvm.script.parser import tirx as T
|
|
|
|
|
|
def test_simple():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
gv = R.sum(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
gv: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
R.output(gv, x_adjoint_out)
|
|
return (gv, (x_adjoint_out,))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
gv: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_assign_binding():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = x
|
|
lv2 = lv1
|
|
gv = R.sum(lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = x
|
|
lv2: R.Tensor((3, 3), dtype="float32") = lv1
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = lv2_adjoint
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
R.output(gv, x_adjoint_out)
|
|
return (gv, (x_adjoint_out,))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = x
|
|
lv2: R.Tensor((3, 3), dtype="float32") = lv1
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_multiple_uses():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, x)
|
|
lv2 = R.add(lv1, x)
|
|
gv = R.sum(lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, x)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.add(lv1, x)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = lv2_adjoint
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv2_adjoint
|
|
x_adjoint1: R.Tensor((3, 3), dtype="float32") = R.add(x_adjoint, lv1_adjoint)
|
|
x_adjoint2: R.Tensor((3, 3), dtype="float32") = R.add(x_adjoint1, lv1_adjoint)
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint2
|
|
R.output(gv, x_adjoint_out)
|
|
return (gv, (x_adjoint_out,))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, x)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.add(lv1, x)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_unused():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, x)
|
|
lv2 = R.add(lv1, x)
|
|
gv = R.sum(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
gv: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, x)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.add(lv1, x)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_default_require_grads():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32"), z: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, y)
|
|
lv2 = R.add(lv1, z)
|
|
gv = R.sum(lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected1:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.add(lv1, z)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = lv2_adjoint
|
|
z_adjoint: R.Tensor((3, 3), dtype="float32") = lv2_adjoint
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
z_adjoint_out: R.Tensor((3, 3), dtype="float32") = z_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out, z_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out, z_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.add(lv1, z)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After1 = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After1, Expected1)
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.add(lv1, z)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = lv2_adjoint
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
R.output(gv, x_adjoint_out)
|
|
return (gv, (x_adjoint_out,))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.add(lv1, z)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After2 = relax.transform.Gradient("main", require_grads=Before["main"].params[0])(Before)
|
|
assert_structural_equal(After2, Expected2)
|
|
|
|
|
|
def test_target_index():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = x
|
|
lv2 = R.sum(x)
|
|
lv3 = R.sum(y)
|
|
R.output(lv1, lv2, lv3)
|
|
return (lv1, lv2, lv3)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((), dtype="float32"), R.Tensor((), dtype="float32")), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = x
|
|
lv2: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
lv3: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
|
|
lv3_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(lv3_adjoint, R.shape([3, 3]))
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
R.output(lv1, lv2, lv3, x_adjoint_out, y_adjoint_out)
|
|
return ((lv1, lv2, lv3), (x_adjoint_out, y_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((), dtype="float32"), R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = x
|
|
lv2: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
lv3: R.Tensor((), dtype="float32") = R.sum(y, axis=None, keepdims=False)
|
|
R.output(lv1, lv2, lv3)
|
|
return (lv1, lv2, lv3)
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main", target_index=2)(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_intermediate_var_require_grads():
|
|
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
|
y = relax.Var("y", R.Tensor((3, 3), "float32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("main", [x, y]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(x * x)
|
|
lv1 = bb.emit(lv0 * y)
|
|
lv2 = bb.emit(lv1 * y)
|
|
gv0 = bb.emit_output(relax.op.sum(lv2))
|
|
bb.emit_func_output(gv0)
|
|
|
|
Before = bb.get()
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"), R.Tensor((), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.multiply(x, x)
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.multiply(lv, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.multiply(lv1, y)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = R.multiply(lv2_adjoint, y)
|
|
lv_adjoint: R.Tensor((3, 3), dtype="float32") = R.multiply(lv1_adjoint, y)
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = R.multiply(lv_adjoint, x)
|
|
lv1_1: R.Tensor((3, 3), dtype="float32") = R.multiply(lv_adjoint, x)
|
|
x_adjoint1: R.Tensor((3, 3), dtype="float32") = R.add(x_adjoint, lv1_1)
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint1
|
|
lv1_adjoint_out: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
|
|
gv_adjoint_out: R.Tensor((), dtype="float32") = gv_adjoint
|
|
R.output(gv, x_adjoint_out, lv1_adjoint_out, gv_adjoint_out)
|
|
return (gv, (x_adjoint_out, lv1_adjoint_out, gv_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.multiply(x, x)
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.multiply(lv, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.multiply(lv1, y)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main", [x, lv1, gv0])(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
# z does not occur in function
|
|
z = relax.Var("z", R.Tensor((3, 3), "float32"))
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main", [x, lv1, z])(Before)
|
|
|
|
|
|
def test_tuple():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3, 3), "float32")),
|
|
y: R.Tensor((3, 3), "float32"),
|
|
z: R.Tensor((3, 3), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv1 = (y, z)
|
|
lv2 = x[0]
|
|
lv3 = lv1[0]
|
|
lv4 = R.add(lv2, lv3)
|
|
gv = R.sum(lv4)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (y, z)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = x[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = lv1[0]
|
|
lv4: R.Tensor((3, 3), dtype="float32") = R.add(lv2, lv3)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv4, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv4_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = lv4_adjoint
|
|
lv3_adjoint: R.Tensor((3, 3), dtype="float32") = lv4_adjoint
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv1_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv3_adjoint, lv)
|
|
lv1_1: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
x_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv2_adjoint, lv1_1)
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint[0]
|
|
z_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint[1]
|
|
x_adjoint_out: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = x_adjoint
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
z_adjoint_out: R.Tensor((3, 3), dtype="float32") = z_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out, z_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out, z_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (y, z)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = x[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = lv1[0]
|
|
lv4: R.Tensor((3, 3), dtype="float32") = R.add(lv2, lv3)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv4, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_tuple_assignment():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = (x, y)
|
|
lv2 = lv1[0]
|
|
lv3 = R.add(lv2, x)
|
|
lv4 = lv1
|
|
lv5 = lv4[0]
|
|
lv6 = R.add(lv5, lv3)
|
|
gv = R.sum(lv6)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = lv1[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.add(lv2, x)
|
|
lv4: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv1
|
|
lv5: R.Tensor((3, 3), dtype="float32") = lv4[0]
|
|
lv6: R.Tensor((3, 3), dtype="float32") = R.add(lv5, lv3)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv6, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv6_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv5_adjoint: R.Tensor((3, 3), dtype="float32") = lv6_adjoint
|
|
lv3_adjoint: R.Tensor((3, 3), dtype="float32") = lv6_adjoint
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv4_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv5_adjoint, lv)
|
|
lv1_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv4_adjoint
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = lv3_adjoint
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv3_adjoint
|
|
lv1_1: R.Tensor((3, 3), dtype="float32") = lv1_adjoint[0]
|
|
lv2_1: R.Tensor((3, 3), dtype="float32") = R.add(lv1_1, lv2_adjoint)
|
|
lv3_1: R.Tensor((3, 3), dtype="float32") = lv1_adjoint[1]
|
|
lv1_adjoint1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv2_1, lv3_1)
|
|
lv4_1: R.Tensor((3, 3), dtype="float32") = lv1_adjoint1[0]
|
|
x_adjoint1: R.Tensor((3, 3), dtype="float32") = R.add(x_adjoint, lv4_1)
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint1[1]
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint1
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = lv1[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.add(lv2, x)
|
|
lv4: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv1
|
|
lv5: R.Tensor((3, 3), dtype="float32") = lv4[0]
|
|
lv6: R.Tensor((3, 3), dtype="float32") = R.add(lv5, lv3)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv6, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_tuple_nested():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tuple(R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3, 3), "float32")), R.Tensor((3, 3), "float32")),
|
|
y: R.Tensor((3, 3), "float32"),
|
|
z: R.Tensor((3, 3), "float32"),
|
|
u: R.Tensor((3, 3), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv1 = ((y, z), u)
|
|
lv2 = x[0]
|
|
lv3 = lv2[0]
|
|
lv4 = lv1[0]
|
|
lv5 = lv4[1]
|
|
lv6 = R.add(lv3, lv5)
|
|
lv7 = x[1]
|
|
lv8 = R.add(lv6, lv7)
|
|
gv = R.sum(lv8)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32"), u: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")) = ((y, z), u)
|
|
lv2: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = x[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = lv2[0]
|
|
lv4: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv1[0]
|
|
lv5: R.Tensor((3, 3), dtype="float32") = lv4[1]
|
|
lv6: R.Tensor((3, 3), dtype="float32") = R.add(lv3, lv5)
|
|
lv7: R.Tensor((3, 3), dtype="float32") = x[1]
|
|
lv8: R.Tensor((3, 3), dtype="float32") = R.add(lv6, lv7)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv8, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv8_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv6_adjoint: R.Tensor((3, 3), dtype="float32") = lv8_adjoint
|
|
lv7_adjoint: R.Tensor((3, 3), dtype="float32") = lv8_adjoint
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv1_1: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
x_adjoint: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")) = ((lv, lv1_1), lv7_adjoint)
|
|
lv3_adjoint: R.Tensor((3, 3), dtype="float32") = lv6_adjoint
|
|
lv5_adjoint: R.Tensor((3, 3), dtype="float32") = lv6_adjoint
|
|
lv2_1: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv4_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv2_1, lv5_adjoint)
|
|
lv3_1: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv1_adjoint: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")) = (lv4_adjoint, lv3_1)
|
|
lv4_1: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv2_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv3_adjoint, lv4_1)
|
|
lv5_1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = x_adjoint[0]
|
|
lv6_1: R.Tensor((3, 3), dtype="float32") = lv5_1[0]
|
|
lv7_1: R.Tensor((3, 3), dtype="float32") = lv2_adjoint[0]
|
|
lv8_1: R.Tensor((3, 3), dtype="float32") = R.add(lv6_1, lv7_1)
|
|
lv9: R.Tensor((3, 3), dtype="float32") = lv5_1[1]
|
|
lv10: R.Tensor((3, 3), dtype="float32") = lv2_adjoint[1]
|
|
lv11: R.Tensor((3, 3), dtype="float32") = R.add(lv9, lv10)
|
|
lv12: R.Tensor((3, 3), dtype="float32") = x_adjoint[1]
|
|
x_adjoint1: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")) = ((lv8_1, lv11), lv12)
|
|
lv13: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv1_adjoint[0]
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = lv13[0]
|
|
z_adjoint: R.Tensor((3, 3), dtype="float32") = lv13[1]
|
|
u_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint[1]
|
|
x_adjoint_out: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")) = x_adjoint1
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
z_adjoint_out: R.Tensor((3, 3), dtype="float32") = z_adjoint
|
|
u_adjoint_out: R.Tensor((3, 3), dtype="float32") = u_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out, z_adjoint_out, u_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out, z_adjoint_out, u_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")), y: R.Tensor((3, 3), dtype="float32"), z: R.Tensor((3, 3), dtype="float32"), u: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), R.Tensor((3, 3), dtype="float32")) = ((y, z), u)
|
|
lv2: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = x[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = lv2[0]
|
|
lv4: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv1[0]
|
|
lv5: R.Tensor((3, 3), dtype="float32") = lv4[1]
|
|
lv6: R.Tensor((3, 3), dtype="float32") = R.add(lv3, lv5)
|
|
lv7: R.Tensor((3, 3), dtype="float32") = x[1]
|
|
lv8: R.Tensor((3, 3), dtype="float32") = R.add(lv6, lv7)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv8, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_tuple_update():
|
|
"""One tensor `x` is used in and out of tuple many times."""
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv0 = (x, y)
|
|
lv1 = R.add(x, y)
|
|
lv2 = lv0[0]
|
|
lv3 = R.add(lv2, y)
|
|
lv4 = R.add(lv1, lv3)
|
|
lv5 = (x, y)
|
|
lv6 = lv5[0]
|
|
lv7 = lv0[0]
|
|
lv8 = R.add(lv4, lv6)
|
|
lv9 = R.add(lv8, lv7)
|
|
gv = R.sum(lv9)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv0: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (x, y)
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = lv0[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.add(lv2, y)
|
|
lv4: R.Tensor((3, 3), dtype="float32") = R.add(lv1, lv3)
|
|
lv5: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (x, y)
|
|
lv6: R.Tensor((3, 3), dtype="float32") = lv5[0]
|
|
lv7: R.Tensor((3, 3), dtype="float32") = lv0[0]
|
|
lv8: R.Tensor((3, 3), dtype="float32") = R.add(lv4, lv6)
|
|
lv9: R.Tensor((3, 3), dtype="float32") = R.add(lv8, lv7)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv9, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv9_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv8_adjoint: R.Tensor((3, 3), dtype="float32") = lv9_adjoint
|
|
lv7_adjoint: R.Tensor((3, 3), dtype="float32") = lv9_adjoint
|
|
lv4_adjoint: R.Tensor((3, 3), dtype="float32") = lv8_adjoint
|
|
lv6_adjoint: R.Tensor((3, 3), dtype="float32") = lv8_adjoint
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv0_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv7_adjoint, lv)
|
|
lv1_1: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv5_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv6_adjoint, lv1_1)
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv5_adjoint[0]
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = lv5_adjoint[1]
|
|
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = lv4_adjoint
|
|
lv3_adjoint: R.Tensor((3, 3), dtype="float32") = lv4_adjoint
|
|
lv2_adjoint: R.Tensor((3, 3), dtype="float32") = lv3_adjoint
|
|
y_adjoint1: R.Tensor((3, 3), dtype="float32") = R.add(y_adjoint, lv3_adjoint)
|
|
lv2_1: R.Tensor((3, 3), dtype="float32") = lv0_adjoint[0]
|
|
lv3_1: R.Tensor((3, 3), dtype="float32") = R.add(lv2_1, lv2_adjoint)
|
|
lv4_1: R.Tensor((3, 3), dtype="float32") = lv0_adjoint[1]
|
|
lv0_adjoint1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv3_1, lv4_1)
|
|
x_adjoint1: R.Tensor((3, 3), dtype="float32") = R.add(x_adjoint, lv1_adjoint)
|
|
y_adjoint2: R.Tensor((3, 3), dtype="float32") = R.add(y_adjoint1, lv1_adjoint)
|
|
lv5_1: R.Tensor((3, 3), dtype="float32") = lv0_adjoint1[0]
|
|
x_adjoint2: R.Tensor((3, 3), dtype="float32") = R.add(x_adjoint1, lv5_1)
|
|
lv6_1: R.Tensor((3, 3), dtype="float32") = lv0_adjoint1[1]
|
|
y_adjoint3: R.Tensor((3, 3), dtype="float32") = R.add(y_adjoint2, lv6_1)
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint2
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint3
|
|
R.output(gv, x_adjoint_out, y_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv0: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (x, y)
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = lv0[0]
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.add(lv2, y)
|
|
lv4: R.Tensor((3, 3), dtype="float32") = R.add(lv1, lv3)
|
|
lv5: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (x, y)
|
|
lv6: R.Tensor((3, 3), dtype="float32") = lv5[0]
|
|
lv7: R.Tensor((3, 3), dtype="float32") = lv0[0]
|
|
lv8: R.Tensor((3, 3), dtype="float32") = R.add(lv4, lv6)
|
|
lv9: R.Tensor((3, 3), dtype="float32") = R.add(lv8, lv7)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv9, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_tuple_op_simple():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((6,), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.split(x, 2)
|
|
lv2 = R.concat(lv1)
|
|
gv = R.sum(lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((6,), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((6,), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(x, indices_or_sections=2, axis=0)
|
|
lv2: R.Tensor((6,), dtype="float32") = R.concat(lv1, axis=0)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv2_adjoint: R.Tensor((6,), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([6]))
|
|
lv1_adjoint: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(lv2_adjoint, indices_or_sections=[3], axis=0)
|
|
x_adjoint: R.Tensor((6,), dtype="float32") = R.concat(lv1_adjoint, axis=0)
|
|
x_adjoint_out: R.Tensor((6,), dtype="float32") = x_adjoint
|
|
R.output(gv, x_adjoint_out)
|
|
return (gv, (x_adjoint_out,))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((6,), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(x, indices_or_sections=2, axis=0)
|
|
lv2: R.Tensor((6,), dtype="float32") = R.concat(lv1, axis=0)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv2, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_tuple_op_construct():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3,), "float32"), y: R.Tuple(R.Tensor((3, ), "float32"), R.Tensor((3, ), "float32")),):
|
|
with R.dataflow():
|
|
lv1 = (x, x)
|
|
lv2 = R.concat(lv1)
|
|
lv3 = R.concat((x, x))
|
|
lv4 = R.concat(y)
|
|
lv5 = R.add(lv2, lv3)
|
|
lv6 = R.add(lv5, lv4)
|
|
gv = R.sum(lv6)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3,), dtype="float32"), y: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32"))) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3,), dtype="float32"), R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")))):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = (x, x)
|
|
lv2: R.Tensor((6,), dtype="float32") = R.concat(lv1, axis=0)
|
|
lv3: R.Tensor((6,), dtype="float32") = R.concat((x, x), axis=0)
|
|
lv4: R.Tensor((6,), dtype="float32") = R.concat(y, axis=0)
|
|
lv5: R.Tensor((6,), dtype="float32") = R.add(lv2, lv3)
|
|
lv6: R.Tensor((6,), dtype="float32") = R.add(lv5, lv4)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv6, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv6_adjoint: R.Tensor((6,), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([6]))
|
|
lv5_adjoint: R.Tensor((6,), dtype="float32") = lv6_adjoint
|
|
lv4_adjoint: R.Tensor((6,), dtype="float32") = lv6_adjoint
|
|
lv2_adjoint: R.Tensor((6,), dtype="float32") = lv5_adjoint
|
|
lv3_adjoint: R.Tensor((6,), dtype="float32") = lv5_adjoint
|
|
y_adjoint: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(lv4_adjoint, indices_or_sections=[3], axis=0)
|
|
lv: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(lv3_adjoint, indices_or_sections=[3], axis=0)
|
|
x_adjoint: R.Tensor((3,), dtype="float32") = lv[0]
|
|
lv1_1: R.Tensor((3,), dtype="float32") = lv[1]
|
|
x_adjoint1: R.Tensor((3,), dtype="float32") = R.add(x_adjoint, lv1_1)
|
|
lv1_adjoint: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(lv2_adjoint, indices_or_sections=[3], axis=0)
|
|
lv2_1: R.Tensor((3,), dtype="float32") = lv1_adjoint[0]
|
|
x_adjoint2: R.Tensor((3,), dtype="float32") = R.add(x_adjoint1, lv2_1)
|
|
lv3_1: R.Tensor((3,), dtype="float32") = lv1_adjoint[1]
|
|
x_adjoint3: R.Tensor((3,), dtype="float32") = R.add(x_adjoint2, lv3_1)
|
|
x_adjoint_out: R.Tensor((3,), dtype="float32") = x_adjoint3
|
|
y_adjoint_out: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = y_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3,), dtype="float32"), y: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32"))) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = (x, x)
|
|
lv2: R.Tensor((6,), dtype="float32") = R.concat(lv1, axis=0)
|
|
lv3: R.Tensor((6,), dtype="float32") = R.concat((x, x), axis=0)
|
|
lv4: R.Tensor((6,), dtype="float32") = R.concat(y, axis=0)
|
|
lv5: R.Tensor((6,), dtype="float32") = R.add(lv2, lv3)
|
|
lv6: R.Tensor((6,), dtype="float32") = R.add(lv5, lv4)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv6, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_tuple_op_const():
|
|
c1 = R.const(np.zeros(3).astype(np.float32))
|
|
c2 = R.const(np.zeros(3).astype(np.float32))
|
|
c3 = R.const(np.zeros(3).astype(np.float32))
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3,), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.concat((c1, c2))
|
|
lv2 = R.concat((c3, x))
|
|
lv3 = R.concat((x, x))
|
|
lv4 = R.add(lv1, lv2)
|
|
lv5 = R.add(lv4, lv3)
|
|
gv = R.sum(lv5)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3,), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3,), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((6,), dtype="float32") = R.concat((c1, c2), axis=0)
|
|
lv2: R.Tensor((6,), dtype="float32") = R.concat((c3, x), axis=0)
|
|
lv3: R.Tensor((6,), dtype="float32") = R.concat((x, x), axis=0)
|
|
lv4: R.Tensor((6,), dtype="float32") = R.add(lv1, lv2)
|
|
lv5: R.Tensor((6,), dtype="float32") = R.add(lv4, lv3)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv5, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv5_adjoint: R.Tensor((6,), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([6]))
|
|
lv4_adjoint: R.Tensor((6,), dtype="float32") = lv5_adjoint
|
|
lv3_adjoint: R.Tensor((6,), dtype="float32") = lv5_adjoint
|
|
lv2_adjoint: R.Tensor((6,), dtype="float32") = lv4_adjoint
|
|
lv: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(lv3_adjoint, indices_or_sections=[3], axis=0)
|
|
x_adjoint: R.Tensor((3,), dtype="float32") = lv[0]
|
|
lv1_1: R.Tensor((3,), dtype="float32") = lv[1]
|
|
x_adjoint1: R.Tensor((3,), dtype="float32") = R.add(x_adjoint, lv1_1)
|
|
lv2_1: R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")) = R.split(lv2_adjoint, indices_or_sections=[3], axis=0)
|
|
lv3_1: R.Tensor((3,), dtype="float32") = lv2_1[1]
|
|
x_adjoint2: R.Tensor((3,), dtype="float32") = R.add(x_adjoint1, lv3_1)
|
|
x_adjoint_out: R.Tensor((3,), dtype="float32") = x_adjoint2
|
|
R.output(gv, x_adjoint_out)
|
|
return (gv, (x_adjoint_out,))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3,), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((6,), dtype="float32") = R.concat((c1, c2), axis=0)
|
|
lv2: R.Tensor((6,), dtype="float32") = R.concat((c3, x), axis=0)
|
|
lv3: R.Tensor((6,), dtype="float32") = R.concat((x, x), axis=0)
|
|
lv4: R.Tensor((6,), dtype="float32") = R.add(lv1, lv2)
|
|
lv5: R.Tensor((6,), dtype="float32") = R.add(lv4, lv3)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv5, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After["main_adjoint"], Expected["main_adjoint"])
|
|
|
|
|
|
def test_const():
|
|
"""const could be used in variable assignment, call argument, and as a part of tuple"""
|
|
cst = relax.const(np.ones((3, 3)), "float32")
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, cst)
|
|
lv2 = cst
|
|
lv3 = (cst, (cst, lv1))
|
|
lv4 = lv3[1]
|
|
lv5 = lv4[1]
|
|
lv6 = R.subtract(lv5, lv2)
|
|
gv = R.sum(lv6)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, cst)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = cst
|
|
lv3: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))) = (cst, (cst, lv1))
|
|
lv4: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv3[1]
|
|
lv5: R.Tensor((3, 3), dtype="float32") = lv4[1]
|
|
lv6: R.Tensor((3, 3), dtype="float32") = R.subtract(lv5, lv2)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv6, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv6_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv5_adjoint: R.Tensor((3, 3), dtype="float32") = lv6_adjoint
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv4_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = (lv, lv5_adjoint)
|
|
lv1_1: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
lv3_adjoint: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))) = (lv1_1, lv4_adjoint)
|
|
lv2_1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv3_adjoint[1]
|
|
lv1_adjoint: R.Tensor((3, 3), dtype="float32") = lv2_1[1]
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = R.zeros(R.shape([3, 3]), dtype="float32")
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, cst)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = cst
|
|
lv3: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))) = (cst, (cst, lv1))
|
|
lv4: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")) = lv3[1]
|
|
lv5: R.Tensor((3, 3), dtype="float32") = lv4[1]
|
|
lv6: R.Tensor((3, 3), dtype="float32") = R.subtract(lv5, lv2)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv6, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_simplify_matmul_pattern():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.permute_dims(x)
|
|
lv2 = R.permute_dims(y)
|
|
lv3 = R.matmul(lv1, lv2, out_dtype="float32")
|
|
gv = R.sum(lv3)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.permute_dims(x, axes=None)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.permute_dims(y, axes=None)
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.matmul(lv1, lv2, out_dtype="float32")
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv3, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv3_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 3]))
|
|
lv: R.Tensor((3, 3), dtype="float32") = R.permute_dims(lv3_adjoint, axes=[1, 0])
|
|
lv1_1: R.Tensor((3, 3), dtype="float32") = R.permute_dims(x, axes=[1, 0])
|
|
y_adjoint: R.Tensor((3, 3), dtype="float32") = R.matmul(lv, lv1_1)
|
|
lv2_1: R.Tensor((3, 3), dtype="float32") = R.permute_dims(y, axes=[1, 0])
|
|
lv3_1: R.Tensor((3, 3), dtype="float32") = R.permute_dims(lv3_adjoint, axes=[1, 0])
|
|
x_adjoint: R.Tensor((3, 3), dtype="float32") = R.matmul(lv2_1, lv3_1)
|
|
x_adjoint_out: R.Tensor((3, 3), dtype="float32") = x_adjoint
|
|
y_adjoint_out: R.Tensor((3, 3), dtype="float32") = y_adjoint
|
|
R.output(gv, x_adjoint_out, y_adjoint_out)
|
|
return (gv, (x_adjoint_out, y_adjoint_out))
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.permute_dims(x, axes=None)
|
|
lv2: R.Tensor((3, 3), dtype="float32") = R.permute_dims(y, axes=None)
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.matmul(lv1, lv2, out_dtype="float32")
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv3, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_shape_expr():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 4), "float32")):
|
|
with R.dataflow():
|
|
s = R.shape([3, 2, 2])
|
|
lv = R.reshape(x, s)
|
|
gv = R.sum(lv)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 4), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 4), dtype="float32"))):
|
|
with R.dataflow():
|
|
s: R.Shape([3, 2, 2]) = R.shape([3, 2, 2])
|
|
lv: R.Tensor((3, 2, 2), dtype="float32") = R.reshape(x, s)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
|
|
gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv_adjoint: R.Tensor((3, 2, 2), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([3, 2, 2]))
|
|
x_adjoint: R.Tensor((3, 4), dtype="float32") = R.reshape(lv_adjoint, R.shape([3, 4]))
|
|
x_adjoint_out: R.Tensor((3, 4), dtype="float32") = x_adjoint
|
|
R.output(gv, x_adjoint_out)
|
|
return (gv, (x_adjoint_out,))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 4), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
s: R.Shape([3, 2, 2]) = R.shape([3, 2, 2])
|
|
lv: R.Tensor((3, 2, 2), dtype="float32") = R.reshape(x, s)
|
|
gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
|
|
R.output(gv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_params_copy():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x0: R.Tensor((3, 3), "float32"),
|
|
x1: R.Tensor((3, 3), "float32"),
|
|
x2: R.Tensor((3, 3), "float32"),
|
|
x3: R.Tensor((3, 3), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.add(x0, x1)
|
|
lv1 = R.add(x2, x3)
|
|
lv2 = R.add(lv0, lv1)
|
|
gv = R.sum(lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
assert len(Before["main"].params) == len(After["main"].params)
|
|
assert len(Before["main"].params) == len(After["main_adjoint"].params)
|
|
for i in range(len(After["main"].params)):
|
|
assert Before["main"].params[i] == After["main"].params[i]
|
|
assert Before["main"].params[i] != After["main_adjoint"].params[i]
|
|
|
|
|
|
def test_function_copy():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x0: R.Tensor((3, 3), "float32"),
|
|
x1: R.Tensor((3, 3), "float32"),
|
|
x2: R.Tensor((3, 3), "float32"),
|
|
x3: R.Tensor((3, 3), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.add(x0, x1)
|
|
lv1 = R.add(x2, x3)
|
|
lv2 = R.add(lv0, lv1)
|
|
gv = R.sum(lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
|
|
# After should have the same "main" function as Before
|
|
assert_structural_equal(Before["main"], After["main"])
|
|
|
|
# the first bindings of After["main_adjoint"] should be the same as Before["main"]
|
|
old_bindings = Before["main"].body.blocks[0].bindings
|
|
old_bindings_len = len(old_bindings)
|
|
new_bindings = After["main_adjoint"].body.blocks[0].bindings[:old_bindings_len]
|
|
assert_structural_equal(old_bindings, new_bindings, True)
|
|
relax.analysis.well_formed(After)
|
|
|
|
|
|
def test_tir_copy():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x0: R.Tensor(("n", "n"), "float32"),
|
|
x1: R.Tensor(("n", "n"), "float32"),
|
|
x2: R.Tensor(("n", "n"), "float32"),
|
|
x3: R.Tensor(("n", "n"), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.add(x0, x1)
|
|
lv1 = R.add(x2, x3)
|
|
lv2 = R.add(lv0, lv1)
|
|
gv = R.sum(lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
relax.analysis.well_formed(After)
|
|
|
|
|
|
def test_report_error():
|
|
@I.ir_module(s_tir=True)
|
|
class TargetNotTensor:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.sum(x)
|
|
gv = R.tuple(lv1, lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main")(TargetNotTensor)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class TargetNotScalar:
|
|
@R.function
|
|
def main(x0: R.Tensor((3, 3), "float32"), x1: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
gv = R.add(x0, x1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main")(TargetNotScalar)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class TargetNotFloat:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
gv = R.const(1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main")(TargetNotFloat)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class ReturnScalarAndWrongTargetIndex:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
gv = R.sum(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main", target_index=1)(ReturnScalarAndWrongTargetIndex)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class ReturnTupleAndWrongTargetIndex:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
gv1 = R.sum(x)
|
|
gv2 = R.sum(y)
|
|
R.output(gv1, gv2)
|
|
return gv1, gv2
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main", target_index=2)(ReturnTupleAndWrongTargetIndex)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class IndexedTargetNotVar:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
gv = R.sum(x)
|
|
R.output(gv)
|
|
return gv, (gv, gv)
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main", target_index=1)(IndexedTargetNotVar)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class NoDataflow:
|
|
@R.function
|
|
def main(x0: R.Tensor((3, 3), "float32")):
|
|
gv = R.sum(x0)
|
|
return gv
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main")(NoDataflow)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class MultiBlocks:
|
|
@R.function
|
|
def main(x0: R.Tensor((3, 3), "float32"), x1: R.Tensor((3, 3), "float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
gv = R.add(x0, x1)
|
|
R.output(gv)
|
|
# block 1
|
|
gv1 = R.sum(x0)
|
|
return gv1
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main")(MultiBlocks)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class NormalModule:
|
|
@R.function
|
|
def main(x0: R.Tensor((3, 3), "float32"), x1: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
gv = R.sum(x0)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def sum(
|
|
rxplaceholder: T.Buffer((T.int64(3), T.int64(3)), "float32"),
|
|
rxplaceholder_red: T.Buffer((), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for k0, k1 in T.grid(T.int64(3), T.int64(3)):
|
|
with T.sblock("rxplaceholder_red"):
|
|
v_k0, v_k1 = T.axis.remap("RR", [k0, k1])
|
|
T.reads(rxplaceholder[v_k0, v_k1])
|
|
T.writes(rxplaceholder_red[()])
|
|
with T.init():
|
|
rxplaceholder_red[()] = T.float32(0)
|
|
rxplaceholder_red[()] = rxplaceholder_red[()] + rxplaceholder[v_k0, v_k1]
|
|
|
|
# no such function
|
|
with pytest.raises(ValueError):
|
|
relax.transform.Gradient("main1")(NormalModule)
|
|
# wrong function type
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("sum")(NormalModule)
|
|
# no such var
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main", require_grads=MultiBlocks["main"].params[0])(NormalModule)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class IntDtype:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "int64")):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, x)
|
|
gv = R.sum(lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main")(IntDtype)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class IntDtypeTuple:
|
|
@R.function
|
|
def main(x: R.Tuple(R.Tensor((3, 3), "int64"), R.Tensor((3, 3), "int64"))):
|
|
with R.dataflow():
|
|
lv1 = x[0]
|
|
lv2 = x[1]
|
|
lv3 = R.add(lv1, lv2)
|
|
gv = R.sum(lv3)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.Gradient("main")(IntDtypeTuple)
|
|
|
|
|
|
def test_mlp_script():
|
|
"""
|
|
An example of single layer multi-layer perceptron. You can add extra layers if you want.
|
|
|
|
For n-layer perceptron, see test_transform_gradient_numeric.py.
|
|
"""
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 10), "float32"),
|
|
w0: R.Tensor((10, 5), "float32"),
|
|
b0: R.Tensor((5,), "float32"),
|
|
label: R.Tensor((3, 5), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, w0)
|
|
out = R.add(lv0, b0)
|
|
logits = R.nn.log_softmax(out)
|
|
loss = R.nn.cross_entropy_with_logits(logits, label)
|
|
R.output(loss)
|
|
return loss
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main_adjoint(x: R.Tensor((3, 10), dtype="float32"), w0: R.Tensor((10, 5), dtype="float32"), b0: R.Tensor((5,), dtype="float32"), label: R.Tensor((3, 5), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((10, 5), dtype="float32"), R.Tensor((5,), dtype="float32"))):
|
|
with R.dataflow():
|
|
lv0: R.Tensor((3, 5), dtype="float32") = R.matmul(x, w0)
|
|
out: R.Tensor((3, 5), dtype="float32") = R.add(lv0, b0)
|
|
logits: R.Tensor((3, 5), dtype="float32") = R.nn.log_softmax(out, axis=-1)
|
|
loss: R.Tensor((), dtype="float32") = R.nn.cross_entropy_with_logits(logits, label)
|
|
loss_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
|
|
lv: R.Tensor((), dtype="float32") = R.divide(loss_adjoint, R.const(3, "float32"))
|
|
lv1: R.Tensor((), dtype="float32") = R.negative(lv)
|
|
logits_adjoint: R.Tensor((3, 5), dtype="float32") = R.multiply(lv1, label)
|
|
lv3: R.Tensor((3, 1), dtype="float32") = R.sum(logits_adjoint, axis=[-1], keepdims=True)
|
|
lv4: R.Tensor((3, 5), dtype="float32") = R.exp(logits)
|
|
lv5: R.Tensor((3, 5), dtype="float32") = R.multiply(lv3, lv4)
|
|
out_adjoint: R.Tensor((3, 5), dtype="float32") = R.subtract(logits_adjoint, lv5)
|
|
lv0_adjoint: R.Tensor((3, 5), dtype="float32") = out_adjoint
|
|
b0_adjoint: R.Tensor((5,), dtype="float32") = R.collapse_sum_to(out_adjoint, R.shape([5]))
|
|
lv7: R.Tensor((10, 3), dtype="float32") = R.permute_dims(x, axes=[1, 0])
|
|
w0_adjoint: R.Tensor((10, 5), dtype="float32") = R.matmul(lv7, lv0_adjoint)
|
|
w0_adjoint_out: R.Tensor((10, 5), dtype="float32") = w0_adjoint
|
|
b0_adjoint_out: R.Tensor((5,), dtype="float32") = b0_adjoint
|
|
R.output(loss, w0_adjoint_out, b0_adjoint_out)
|
|
return (loss, (w0_adjoint_out, b0_adjoint_out))
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 10), dtype="float32"), w0: R.Tensor((10, 5), dtype="float32"), b0: R.Tensor((5,), dtype="float32"), label: R.Tensor((3, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
|
|
with R.dataflow():
|
|
lv0: R.Tensor((3, 5), dtype="float32") = R.matmul(x, w0)
|
|
out: R.Tensor((3, 5), dtype="float32") = R.add(lv0, b0)
|
|
logits: R.Tensor((3, 5), dtype="float32") = R.nn.log_softmax(out, axis=-1)
|
|
loss: R.Tensor((), dtype="float32") = R.nn.cross_entropy_with_logits(logits, label)
|
|
R.output(loss)
|
|
return loss
|
|
# fmt: on
|
|
|
|
After = relax.transform.Gradient("main", require_grads=Before["main"].params[1:3])(Before)
|
|
assert_structural_equal(After, Expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|