/* * 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. */ /*! * \file tvm/te/autodiff.h * \brief Automatic differentiation of tensor expressions. */ #ifndef TVM_TE_AUTODIFF_H_ #define TVM_TE_AUTODIFF_H_ #include <tvm/runtime/object.h> #include <tvm/tir/expr.h> #include "tensor.h" namespace tvm { /*! \brief Tensor expression language DSL. */ namespace te { /*! * \brief Take the derivative of the expression with respect to the given variable. * \param expr The expression to differentiate. * \param var The variable to differentiate with respect to. * \return The expression for the derivative. */ PrimExpr Derivative(const PrimExpr& expr, const Var& var); /*! * \brief Get the tensor representing the Jacobian of the output with respect to the input. * * Note that if \p output depends on \p input indirectly (by using some other tensor * depending on \p input), this dependency won't contribute to the resulting Jacobian. * For such cases use the function ::Gradient. * * \param output The tensor to differentiate. * \param input The input tensor, which \p output should directly use. * \return The tensor representing the Jacobian of shape `output.shape + input.shape`. */ Tensor Jacobian(const Tensor& output, const Tensor& input); /*! * \brief The building block for reverse-mode AD. * * Differentiate \p output wrt \p input and multiply the result by \p head on the left using tensor * dot product. \p input must be an immediate dependency of \p output (must be called from within * the body of \p output). That is, the function will compute one summand of the adjoint for \p input * given the adjoint for \p output (which is called \p head here). * * \param output The tensor to differentiate. * \param input The input tensor, which \p output should directly use. * \param head The adjoint of \p output. Must be of shape `prefix + output.shape` * \return The tensor of shape `prefix + input.shape` * representing the partial adjoint of \p input wrt one of its consumers (output) */ Tensor VectorJacobianProduct(const Tensor &output, const Tensor &input, const Tensor &head); /*! * \brief Perform reverse mode automatic differentiation. * * Each item of the `result` field of the result is an adjoint for the corresponding item of * \p inputs, i.e. \p head multiplied by the Jacobian of \p output with respect to the * corresponding item of \p inputs. * * \param output The tensor to differentiate. * \param inputs The array of input tensors. When the array is empty, will perform differentiation * wrt all tensors the output depends on. * \param head The adjoint of the output, in other words, some tensor, by which the Jacobians * will be multiplied (using tensordot axes=`output.shape`). * Its shape must be of the form `prefix + output.shape`. If the null pointer is provided, * the identity tensor of shape `output.shape + output.shape` will be used. * \return An array of adjoints corresponding to \p inputs. */ TVM_DLL Array<Tensor> Gradient( const Tensor& output, const Array<Tensor>& inputs, const Tensor& head = Tensor()); } // namespace te } // namespace tvm #endif // TVM_TE_AUTODIFF_H_