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
33namespace tvm {
34/*! \brief Tensor expression language DSL. */
35namespace 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 */
43PrimExpr 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 */
56Tensor 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 */
72Tensor 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 */
90TVM_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