/*
 * 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_