1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file tvm/te/autodiff.h |
22 | * \brief Automatic differentiation of tensor expressions. |
23 | */ |
24 | |
25 | #ifndef TVM_TE_AUTODIFF_H_ |
26 | #define TVM_TE_AUTODIFF_H_ |
27 | |
28 | #include <tvm/runtime/object.h> |
29 | #include <tvm/tir/expr.h> |
30 | |
31 | #include "tensor.h" |
32 | |
33 | namespace tvm { |
34 | /*! \brief Tensor expression language DSL. */ |
35 | namespace te { |
36 | |
37 | /*! |
38 | * \brief Take the derivative of the expression with respect to the given variable. |
39 | * \param expr The expression to differentiate. |
40 | * \param var The variable to differentiate with respect to. |
41 | * \return The expression for the derivative. |
42 | */ |
43 | PrimExpr Derivative(const PrimExpr& expr, const Var& var); |
44 | |
45 | /*! |
46 | * \brief Get the tensor representing the Jacobian of the output with respect to the input. |
47 | * |
48 | * Note that if \p output depends on \p input indirectly (by using some other tensor |
49 | * depending on \p input), this dependency won't contribute to the resulting Jacobian. |
50 | * For such cases use the function ::Gradient. |
51 | * |
52 | * \param output The tensor to differentiate. |
53 | * \param input The input tensor, which \p output should directly use. |
54 | * \return The tensor representing the Jacobian of shape `output.shape + input.shape`. |
55 | */ |
56 | Tensor Jacobian(const Tensor& output, const Tensor& input); |
57 | |
58 | /*! |
59 | * \brief The building block for reverse-mode AD. |
60 | * |
61 | * Differentiate \p output wrt \p input and multiply the result by \p head on the left using tensor |
62 | * dot product. \p input must be an immediate dependency of \p output (must be called from within |
63 | * the body of \p output). That is, the function will compute one summand of the adjoint for \p |
64 | * input given the adjoint for \p output (which is called \p head here). |
65 | * |
66 | * \param output The tensor to differentiate. |
67 | * \param input The input tensor, which \p output should directly use. |
68 | * \param head The adjoint of \p output. Must be of shape `prefix + output.shape` |
69 | * \return The tensor of shape `prefix + input.shape` |
70 | * representing the partial adjoint of \p input wrt one of its consumers (output) |
71 | */ |
72 | Tensor VectorJacobianProduct(const Tensor& output, const Tensor& input, const Tensor& head); |
73 | |
74 | /*! |
75 | * \brief Perform reverse mode automatic differentiation. |
76 | * |
77 | * Each item of the `result` field of the result is an adjoint for the corresponding item of |
78 | * \p inputs, i.e. \p head multiplied by the Jacobian of \p output with respect to the |
79 | * corresponding item of \p inputs. |
80 | * |
81 | * \param output The tensor to differentiate. |
82 | * \param inputs The array of input tensors. When the array is empty, will perform differentiation |
83 | * wrt all tensors the output depends on. |
84 | * \param head The adjoint of the output, in other words, some tensor, by which the Jacobians |
85 | * will be multiplied (using tensordot axes=`output.shape`). |
86 | * Its shape must be of the form `prefix + output.shape`. If the null pointer is |
87 | * provided, the identity tensor of shape `output.shape + output.shape` will be used. \return An |
88 | * array of adjoints corresponding to \p inputs. |
89 | */ |
90 | TVM_DLL Array<Tensor> Gradient(const Tensor& output, const Array<Tensor>& inputs, |
91 | const Tensor& head = Tensor()); |
92 | |
93 | } // namespace te |
94 | } // namespace tvm |
95 | |
96 | #endif // TVM_TE_AUTODIFF_H_ |
97 | |