text: /**
   * Copyright 2019 Huawei Technologies Co., Ltd
   *
   * Licensed 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.
   */
  #ifndef PASS_AUTODIFF_H_
  #define PASS_AUTODIFF_H_
  
  #include <tvm/ir.h>
  #include <tvm/tensor.h>
  
  #include <tvm.h>
  
  namespace akg {
  namespace ir {
  class DifferentiationResult;
  
  /*! \brief Node to represent a differentiation result */
  class DifferentiationResultNode : public Node {
   public:
    /*! \brief The requested adjoints, i.e. Jacobians or gradients wrt to the given inputs */
    Array<Tensor> result;
    /*! \brief A map from tensors to the corresponding adjoints (including internal nodes) */
    Map<Tensor, Tensor> adjoints;
    /*! \brief Single summands of the adjoints*/
    Map<Tensor, Map<Tensor, Tensor>> adjoint_summands;
    /*! \brief constructor */
    DifferentiationResultNode() = default;
  
    /*! \brief destructor */
    ~DifferentiationResultNode() = default;
  
    void VisitAttrs(AttrVisitor *v) {
      v->Visit("result", &result);
      v->Visit("adjoints", &adjoints);
      v->Visit("adjoint_summands", &adjoint_summands);
    }
    TVM_DLL static DifferentiationResult make(Array<Tensor> result, Map<Tensor, Tensor> adjoints,
                                              Map<Tensor, Map<Tensor, Tensor>> adjoint_summands);
  
    static constexpr const char *_type_key = "DifferentiationResult";
    TVM_DECLARE_NODE_TYPE_INFO(DifferentiationResultNode, Node);
  };
  
  TVM_DEFINE_NODE_REF(DifferentiationResult, DifferentiationResultNode);
  
  /*! \brief A type of a "local" differentiation function for reverse mode AD
   *
   *  A function of this type is a building block for reverse-mode automatic differentiation. It
   *  should take three tensors: `output`, `input` and `head`, `head` being the adjoint corresponding
   *  to the `output`, and return (a summand of) the adjoint corresponding to the input. In other
   *  words, it should differentiate `output` wrt `input` and multiply the result by `head` with
   *  tensor dot product (`head` should be on the left of the multiplication). `input` should be an
   *  immediate dependency of `output` (should be called from within the body of `output`).
   *
   *  See also ::DiffBuildingBlock, which might be considered the reference implementation.
   */
  using FDiffBuildingBlock =
    std::function<Tensor(const Tensor &output, const Tensor &input, const Tensor &head,
                         const Map<std::string, NodeRef> &attrs, const Array<Tensor> &new_pld_array)>;
  
  /*!
   * \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.
   */
  TVM_DLL Expr Derivative(const Expr &expr, const VarExpr &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 ::Differentiate.
   *
   *  The \p head might be used inside OptimizeReduction() to improve performance
   *
   * \param output The tensor to differentiate.
   * \param input The input tensor, which \p output should directly use.
   * \param used_head Boolean the indicates if the head was already used inside one of the
   *                  optimizations.
   * \param optimize Whether to perform optimizations like lifting of nonzeroness conditions.
   * \param keep_dims Select to keep or not the dimensions equal 1 in newly created Jacobians
   * \param head The adjoint of the output, in other words, some tensor, by which the Jacobians
   *             will be multiplied. 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 The tensor representing the Jacobian of shape `output.shape + input.shape`.
   */
  TVM_DLL Tensor Jacobian(const Tensor &output, const Tensor &input, bool &used_head, bool optimize = true,
                          bool keep_dims = false, const Tensor &head = Tensor());
  
  /*!
   * \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 a 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`
   * \param attrs The additional set of attributes used for the AD algorithms
   * \param new_pld_array The additional array of Tensors used in AD transformations
   * \return The tensor representing the adjoint of \p input of shape `prefix + input.shape`.
   */
  TVM_DLL Tensor DiffBuildingBlock(const Tensor &output, const Tensor &input, const Tensor &head,
                                   const Map<std::string, NodeRef> &attrs, const Array<Tensor> &new_pld_array);
  
  /*!
   * \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. 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.
   * \param attrs The additional set of attributes used for the AD algorithms
   * \param new_pld_array The additional array of Tensors used in AD transformations
   * \param fdiff The function performing differentiation and multiplication, see
   *              ::FDiffBuildingBlock.
   * \param override_deps A map from tensors to their dependencies (`InputTensors()` are used by
   *                      default). Overriding dependencies may be useful to treat a group of tensors
   *                      as a single supertensor. In this case the fdiff functions should also be
   *                      modified accordingly.
   * \return An object of type DifferentiationResult which contains three fields:
   *         - `result` An array of adjoints corresponding to \p inputs.
   *         - `adjoints` A map from tensors to the corresponding adjoints (includes intermediate
   *            tensors).
   *         - `adjoint_summands` A map from tensors to maps from parent tensors to individual
   *            summands of the adjoint.
   */
  TVM_DLL DifferentiationResult
  Differentiate(const Tensor &output, const Array<Tensor> &inputs = Array<Tensor>(), const Tensor &head = Tensor(),
                const Map<std::string, NodeRef> &attrs = {}, const Array<Tensor> &new_pld_array = Array<Tensor>(),
                const FDiffBuildingBlock &fdiff = DiffBuildingBlock,
                const Map<Tensor, Array<Tensor>> &override_deps = Map<Tensor, Array<Tensor>>());
  }  // namespace ir
  }  // namespace akg
  #endif  // PASS_AUTODIFF_H_
  