# 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()