You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/03/22 05:24:55 UTC
[GitHub] [incubator-tvm] yzhliu opened a new pull request #5121: [TE]
reverse-mode autodiff without any optimization
yzhliu opened a new pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121
This is the first PR to bring in previously-implemented tensor-level autodiff.
This PR does not include any optimization, thus produces bad performance. Will submit optimization pass in another two or three PRs, so that not to put too much pressure on reviewers.
Also credit to @sgrechanik-h as I mentioned in the header of each file.
RFC: https://discuss.tvm.ai/t/rfc-bring-in-tensor-expression-autodiff
Please help to review @sgrechanik-h @MarisaKirisame @junrushao1994 @tqchen @hzfan
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen edited a comment on issue #5121: [TE]
reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
tqchen edited a comment on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-606353531
Thanks @yzhliu @sgrechanik-h @MarisaKirisame @hzfan
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] sgrechanik-h commented on issue #5121: [TE]
reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
sgrechanik-h commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-603000753
Thanks for reviving this. The PR looks good to me, but I'm obviously partial.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on issue #5121: [TE] reverse-mode
autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-604099502
CI's green. @tqchen @hzfan check if there's anything needs to be addressed.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396628415
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
+ // may use an arbitrary combiner.
+ // The resulting reduction expression will return a tuple containing
+ // both derivatives and the original results (in exactly this order).
+
+ // Example of a ReduceNode,
+ // reduce(combiner=comm_reducer(result=[(x + y)], lhs=[x], rhs=[y], identity_element=[0f]),
+ // source=[A(k)], axis=[iter_var(k, range(min=0, ext=5))], where=(bool)1, value_index=0)
+
+ // We have to clone the reduction axes because otherwise the original expression
+ // cannot be used together with the derivative (it will lead to errors during lowering)
+ PrimExpr expr_with_new_axes = te::CloneReduction(GetRef<PrimExpr>(op));
+ const ReduceNode* new_op = expr_with_new_axes.as<ReduceNode>();
+
+ // New lhs and rhs variables of the new combiner consist of
+ // variables representing derivatives (which are later derived from new_op->source)
+ // followed by the original variables.
+ Array<Var> new_lhs;
+ for (const auto& var : new_op->combiner->lhs) {
+ new_lhs.push_back(var.copy_with_suffix(".jac"));
+ }
+ for (const auto& var : new_op->combiner->lhs) {
+ new_lhs.push_back(var);
+ }
+
+ Array<Var> new_rhs;
+ for (const auto& var : new_op->combiner->rhs) {
+ new_rhs.push_back(var.copy_with_suffix(".jac"));
+ }
+ for (const auto& var : new_op->combiner->rhs) {
+ new_rhs.push_back(var);
+ }
+
+ // The new combiner result also consists of the resulting derivatives
+ // followed by the original results.
+ Array<PrimExpr> new_result;
+ for (const auto& res : new_op->combiner->result) {
+ // Each resulting derivative is computed as a sum of derivatives
+ // wrt lhs and rhs multiplied by the derivatives of lhs and rhs
+ PrimExpr new_res = make_zero(res.dtype());
+ for (size_t i = 0; i < new_op->combiner->lhs.size(); ++i) {
+ PrimExpr res_di = Derivative(res, new_op->combiner->lhs[i]);
+ // new_lhs[i] is the derivative of lhs[i] (wrt our input tensor)
+ new_res = AddNode::make(new_res, MulNode::make(new_lhs[i], res_di));
+ }
+ for (size_t i = 0; i < new_op->combiner->rhs.size(); ++i) {
+ PrimExpr res_di = Derivative(res, new_op->combiner->rhs[i]);
+ // new_rhs[i] is the derivative of rhs[i] (wrt our input tensor)
+ new_res = AddNode::make(new_res, MulNode::make(new_rhs[i], res_di));
+ }
+ new_result.push_back(new_res);
+ }
+ // add original results
+ for (const auto& res : new_op->combiner->result) {
+ new_result.push_back(res);
+ }
+
+ // The identity is transformed in a similar way
+ Array<PrimExpr> new_identity;
+ for (const auto& id : new_op->combiner->identity_element) {
+ new_identity.push_back(Mutate(id));
+ }
+ for (const auto& id : new_op->combiner->identity_element) {
+ new_identity.push_back(id);
+ }
+
+ // Same as source
+ Array<PrimExpr> new_source;
+ for (const auto& src : new_op->source) {
+ new_source.push_back(Mutate(src));
+ }
+ for (const auto& src : new_op->source) {
+ new_source.push_back(src);
+ }
+
+ CommReducer new_combiner = CommReducerNode::make(new_lhs, new_rhs, new_result, new_identity);
+ // Also simplify the resulting combiner
+ // (mostly to get rid of unused components, e.g., the original expressions)
+ return Simplify(
+ ReduceNode::make(new_combiner, new_source, new_op->axis,
+ new_op->condition, new_op->value_index));
+ }
+
+ PrimExpr VisitExpr_(const CastNode* op) {
+ if (op->dtype.is_float()) {
+ return CastNode::make(op->dtype, Mutate(op->value));
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const NotNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const SelectNode* op) {
+ return SelectNode::make(op->condition,
+ Mutate(op->true_value), Mutate(op->false_value));
+ }
+
+ PrimExpr VisitExpr_(const RampNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const BroadcastNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const ShuffleNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const IntImmNode* op) {
+ return IntImm(op->dtype, 0);
+ }
+
+ PrimExpr VisitExpr_(const FloatImmNode* op) {
+ return FloatImm(op->dtype, 0);
+ }
+
+ PrimExpr VisitExpr_(const StringImmNode* op) NOT_IMPLEMENTED
+
+ private:
+ Tensor input_;
+ Array<PrimExpr> indices_;
+ Var input_var_;
+};
+
+PrimExpr Derivative(const PrimExpr& expr, const Var& var) {
+ return JacobianMutator(var).Mutate(expr);
+}
+
+PrimExpr Jacobian(const PrimExpr& expr, const Tensor& input, const Array<PrimExpr>& indices) {
+ return JacobianMutator(input, indices).Mutate(expr);
+}
+
+Tensor Jacobian(const Tensor& output, const Tensor& input) {
+ const ComputeOpNode* op = output->op.as<ComputeOpNode>();
+ CHECK(op) << "Derivative of this operation is not implemented: " << output->op;
+ bool is_input_tensor = false;
+ for (const Tensor& child : op->InputTensors()) {
+ if (input == child) {
+ is_input_tensor = true;
+ break;
+ }
+ }
+ CHECK(is_input_tensor) << "Jacobian is called on a pair of tensors such that the output "
+ << "does not directly depend on the input.";
+
+ // We have to clone the iteration axes because otherwise the original expression
+ // cannot be used together with the derivative (it will lead to errors during lowering)
+ Array<IterVar> new_axis;
Review comment:
Can we reuse `CloneIterVars` here?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on issue #5121: [TE] reverse-mode
autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
tqchen commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-606353531
Thanks @yzhliu @sgrechanik-h @MarisaKirisame
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] sergei-grechanik commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
sergei-grechanik commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r400596264
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
+ // may use an arbitrary combiner.
+ // The resulting reduction expression will return a tuple containing
+ // both derivatives and the original results (in exactly this order).
Review comment:
Looks more like a bug in lowering of tupled reductions rather than an intended behavior, might deserve a separate bug report.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396814445
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
Review comment:
There are no concrete example. If you dont think it will happend then leave it as is.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] hzfan commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
hzfan commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396631729
##########
File path: python/tvm/te/autodiff.py
##########
@@ -0,0 +1,72 @@
+# 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.
+
+"""
+Automatic differentiation of tensor expressions.
+The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+in [Automatic differentiation for tensor expressions](#2498)
+and [Zero elimination](#2634)
+"""
+from . import _ffi_api
+
+
+def gradient(output, inputs, head=None):
+ """Perform reverse-mode automatic differentiation.
+
+ Parameters
+ ----------
+ output : Tensor
+ The tensor to differentiate.
+
+ inputs : List[Tensor]
+ The list of input tensors to be differentiated wrt.
+
+ head : Tensor
+ 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 `None` is passed, the identity tensor of shape `output.shape + output.shape`
Review comment:
So the default behavior is to return a Jacobian instead of adjoint, right?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396935671
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
+ // may use an arbitrary combiner.
+ // The resulting reduction expression will return a tuple containing
+ // both derivatives and the original results (in exactly this order).
Review comment:
Looking into a bit more, the order actually makes difference. When original init value is different from its derivative init value, and they depends on each other during calculation, we must calculate derivative first (using origin's init value), switch the order in tvm makes the origin value be replaced before using, produces incorrect results.
One example is in the test case,
```python
def fcombine(x, y):
return x*y
def fidentity(t0):
return tvm.tir.const(1, t0)
prod = te.comm_reducer(fcombine, fidentity, name='prod')
B = te.compute((10, 10), lambda i, j: prod(A0[i, k] + A0[k, i], axis=k), name='B')
check_grad(B, A0)
```
Correct result (derivative first):
```
produce B.jacobian {
for (i, 0, 10) {
for (j, 0, 10) {
for (jac_i0, 0, 10) {
for (jac_i1, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 0f
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 1f
for (k, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = ((B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)])) + ((float32(((jac_i0 == i) && (jac_i1 == k))) + float32(((jac_i0 == k) && (jac_i1 == i))))*B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]))
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = (B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)]))
}
}
}
}
}
}
Output B.jacobian.v0
```
Incorrect result (origin first):
```
produce B.jacobian {
for (i, 0, 10) {
for (j, 0, 10) {
for (jac_i0, 0, 10) {
for (jac_i1, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 1f
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 0f
for (k, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = (B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)]))
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = ((B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)])) + ((float32(((jac_i0 == i) && (jac_i1 == k))) + float32(((jac_i0 == k) && (jac_i1 == i))))*B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]))
}
}
}
}
}
}
Output B.jacobian.v1
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396135612
##########
File path: python/tvm/testing.py
##########
@@ -68,6 +69,10 @@ def check_numerical_grads(function, input_values, grad_values, function_value=No
rtol : float, optional
Relative tolerance.
+
+ acceptable_fail_percentage : float, optional
Review comment:
```suggestion
acceptable_fail_percentage : Optional[float]
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen merged pull request #5121: [TE] reverse-mode
autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
tqchen merged pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r400643439
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
+ // may use an arbitrary combiner.
+ // The resulting reduction expression will return a tuple containing
+ // both derivatives and the original results (in exactly this order).
Review comment:
@tqchen can you take a quick look and see if it is a bug in tuple reductions?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396670848
##########
File path: python/tvm/testing.py
##########
@@ -68,6 +69,10 @@ def check_numerical_grads(function, input_values, grad_values, function_value=No
rtol : float, optional
Relative tolerance.
+
+ acceptable_fail_percentage : float, optional
Review comment:
actually this is the problem of numerical derivative in `check_numerical_grads` itself. I can change to assert them ("floor", "ceil", "trunc", "round") to be exact zero.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396814133
##########
File path: python/tvm/testing.py
##########
@@ -68,6 +69,10 @@ def check_numerical_grads(function, input_values, grad_values, function_value=No
rtol : float, optional
Relative tolerance.
+
+ acceptable_fail_percentage : float, optional
Review comment:
I see. both are fine IMO.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on issue #5121: [TE] reverse-mode
autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-605845840
Kindly ping @tqchen , can we merge if it looks good?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on issue #5121: [TE] reverse-mode
autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-603446436
@MarisaKirisame @tqchen @hzfan Could you review again?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on issue #5121: [TE]
reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-603454210
(of course, not in this PR)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396141425
##########
File path: python/tvm/testing.py
##########
@@ -68,6 +69,10 @@ def check_numerical_grads(function, input_values, grad_values, function_value=No
rtol : float, optional
Relative tolerance.
+
+ acceptable_fail_percentage : float, optional
Review comment:
Why is this needed? My first instinct tell me that something else is wrong, and this is merely "hiding the problem".
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on issue #5121: [TE] reverse-mode
autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-603648759
@MarisaKirisame sure, I will try.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on issue #5121: [TE]
reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on issue #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#issuecomment-603454001
@yzhliu can you do forward mode automatic differentiation? It is easy considering you have jacobian - you only need to do JacobianVectorProduct instead of VectorJacbianProduct
It is useful in higher order derivative a la hessian vector product.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396153769
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
Review comment:
Is it possible to have reduce inside other expression as well?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396549865
##########
File path: include/tvm/te/autodiff.h
##########
@@ -0,0 +1,100 @@
+/*
+ * 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.
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
Review comment:
You can add sgrechanik-h as a co-author in the commit message, this will list him as an author as well
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396681870
##########
File path: python/tvm/te/autodiff.py
##########
@@ -0,0 +1,72 @@
+# 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.
+
+"""
+Automatic differentiation of tensor expressions.
+The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+in [Automatic differentiation for tensor expressions](#2498)
+and [Zero elimination](#2634)
+"""
+from . import _ffi_api
+
+
+def gradient(output, inputs, head=None):
+ """Perform reverse-mode automatic differentiation.
+
+ Parameters
+ ----------
+ output : Tensor
+ The tensor to differentiate.
+
+ inputs : List[Tensor]
+ The list of input tensors to be differentiated wrt.
+
+ head : Tensor
+ 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 `None` is passed, the identity tensor of shape `output.shape + output.shape`
Review comment:
that's right. more precisely, that's because the arguments are one output and multiple inputs, instead of one input and multiple outputs. if y is the only output, dy/dx is jacobian, it's also adjoint(x) for the previous layer. it depends on what aspect you want to emphasize, you use different terms.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396140353
##########
File path: include/tvm/te/autodiff.h
##########
@@ -0,0 +1,100 @@
+/*
+ * 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.
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+
+#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 PartialAdjoint(const Tensor& output, const Tensor& input, const Tensor& head);
Review comment:
Can you call it "VectorJacobianProduct", or "VJP"? That's how other ppl called it.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] tqchen commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396549865
##########
File path: include/tvm/te/autodiff.h
##########
@@ -0,0 +1,100 @@
+/*
+ * 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.
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
Review comment:
You can add sgrechanik-h as a co-author in the commit message, this will list him as an author as well(and possibly remove the need of explicit attribution as they are part of the git history)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396935671
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
+ // may use an arbitrary combiner.
+ // The resulting reduction expression will return a tuple containing
+ // both derivatives and the original results (in exactly this order).
Review comment:
Looking into a bit more, the order actually makes difference. When original init value is different from its derivative init value, and they depends on each other during calculation. we must calculate derivative first (using origin's init value), switch the order in tvm makes the origin value be replaced, produces incorrect results.
One example is in the test case,
```python
def fcombine(x, y):
return x*y
def fidentity(t0):
return tvm.tir.const(1, t0)
prod = te.comm_reducer(fcombine, fidentity, name='prod')
B = te.compute((10, 10), lambda i, j: prod(A0[i, k] + A0[k, i], axis=k), name='B')
check_grad(B, A0)
```
Correct result (derivative first):
```
produce B.jacobian {
for (i, 0, 10) {
for (j, 0, 10) {
for (jac_i0, 0, 10) {
for (jac_i1, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 0f
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 1f
for (k, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = ((B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)])) + ((float32(((jac_i0 == i) && (jac_i1 == k))) + float32(((jac_i0 == k) && (jac_i1 == i))))*B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]))
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = (B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)]))
}
}
}
}
}
}
Output B.jacobian.v0
```
Incorrect result (origin first):
```
produce B.jacobian {
for (i, 0, 10) {
for (j, 0, 10) {
for (jac_i0, 0, 10) {
for (jac_i1, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 1f
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = 0f
for (k, 0, 10) {
B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = (B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)]))
B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)] = ((B.jacobian.v1[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]*(A0[((i*10) + k)] + A0[((k*10) + i)])) + ((float32(((jac_i0 == i) && (jac_i1 == k))) + float32(((jac_i0 == k) && (jac_i1 == i))))*B.jacobian.v0[((((i*1000) + (j*100)) + (jac_i0*10)) + jac_i1)]))
}
}
}
}
}
}
Output B.jacobian.v1
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] MarisaKirisame commented on a change in pull
request #5121: [TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
MarisaKirisame commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396145968
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
+ // may use an arbitrary combiner.
+ // The resulting reduction expression will return a tuple containing
+ // both derivatives and the original results (in exactly this order).
Review comment:
can you switch the order? most ad code use original result first, derivatives later.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] yzhliu commented on a change in pull request #5121:
[TE] reverse-mode autodiff without any optimization
Posted by GitBox <gi...@apache.org>.
yzhliu commented on a change in pull request #5121: [TE] reverse-mode autodiff without any optimization
URL: https://github.com/apache/incubator-tvm/pull/5121#discussion_r396785055
##########
File path: src/te/autodiff/jacobian.cc
##########
@@ -0,0 +1,381 @@
+/*
+ * 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 jacobian.cc
+ * \brief Calculate Jacobian of two tensors dY/dX.
+ * X must be direct input tensor of Y.
+ * The result Jacobian shape will be (Y.shape, X.shape)
+ * The algorithm was initially implemented by Sergei Grechanik (sgrechanik-h)
+ * in [Automatic differentiation for tensor expressions](#2498)
+ * and [Zero elimination](#2634)
+ */
+#include <tvm/te/autodiff.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+#include <topi/transform.h>
+#include <memory>
+#include "ad_util.h"
+
+namespace tvm {
+namespace te {
+
+#define NOT_IMPLEMENTED \
+ { LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); throw; }
+
+/*! \brief Differentiate an expression wrt a variable or a tensor element */
+class JacobianMutator : public ExprMutator {
+ public:
+ /*!
+ * \brief Differentiate wrt `input(indices)`.
+ * \param input The input tensor.
+ * \param indices The indices of the element with respect to which to differentiate.
+ */
+ explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
+ : input_(input), indices_(indices) {}
+ /*!
+ * \brief Differentiate wrt the input variable.
+ * \param input The input variable.
+ */
+ explicit JacobianMutator(Var input) : input_var_(input) {}
+
+ PrimExpr Mutate(PrimExpr e) {
+ if (e.dtype().is_int() || e.dtype().is_uint()) {
+ LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
+ << " e = " << e;
+ return make_zero(e.dtype());
+ } else {
+ return ExprMutator::VisitExpr(e);
+ }
+ }
+
+ PrimExpr VisitExpr_(const VarNode* op) {
+ if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
+ return FloatImm(op->dtype, 1.0);
+ } else {
+ return make_zero(op->dtype);
+ }
+ }
+
+ PrimExpr VisitExpr_(const LoadNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const CallNode* op) {
+ PrimExpr expr = GetRef<PrimExpr>(op);
+ if (op->call_type == CallNode::CallType::Halide) {
+ if (input_.get() && op->func.same_as(input_->op) &&
+ op->value_index == input_->value_index) {
+ // Tensor(indices)
+ CHECK_EQ(indices_.size(), op->args.size());
+ PrimExpr condition = const_true();
+ for (size_t i = 0; i < input_.ndim(); ++i) {
+ condition = AndNode::make(condition, EQNode::make(indices_[i], op->args[i]));
+ }
+ return CastNode::make(op->dtype, condition);
+ } else {
+ return make_zero(op->dtype);
+ }
+ } else if (op->call_type == CallNode::CallType::PureIntrinsic) {
+ static std::unordered_set<std::string> piecewise_const = {"floor", "ceil", "trunc", "round"};
+ if (op->name == "exp") {
+ return MulNode::make(Mutate(op->args[0]), expr);
+ } else if (op->name == "log") {
+ return DivNode::make(Mutate(op->args[0]), op->args[0]);
+ } else if (op->name == "sigmoid") {
+ return MulNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, SubNode::make(FloatImm(expr.dtype(), 1.0), expr)));
+ } else if (op->name == "sqrt") {
+ return DivNode::make(Mutate(op->args[0]),
+ MulNode::make(expr, FloatImm(expr.dtype(), 2.0)));
+ } else if (op->name == "tanh") {
+ return MulNode::make(Mutate(op->args[0]),
+ SubNode::make(FloatImm(expr.dtype(), 1.0), MulNode::make(expr, expr)));
+ } else if (op->name == "pow") {
+ auto x = op->args[0], y = op->args[1];
+ return expr * (Mutate(y)*log(x) + Mutate(x)*y/x);
+ } else if (op->name == "fabs") {
+ auto type = op->args[0].dtype();
+ return MulNode::make(Mutate(op->args[0]),
+ SelectNode::make(GENode::make(op->args[0], make_zero(type)),
+ FloatImm(type, 1.0), FloatImm(type, -1.0)));
+ } else if (op->name == intrinsic::tvm_if_then_else) {
+ Array<PrimExpr> new_args = {op->args[0],
+ Mutate(op->args[1]),
+ Mutate(op->args[2])};
+ return CallNode::make(op->dtype, op->name, new_args,
+ op->call_type, op->func, op->value_index);
+ } else if (piecewise_const.count(op->name)) {
+ return FloatImm(expr.dtype(), 0.0);
+ } else {
+ throw dmlc::Error("Derivative of this intrinsic is not implemented: " + op->name);
+ }
+ }
+ NOT_IMPLEMENTED
+ }
+
+ PrimExpr VisitExpr_(const AddNode* op) {
+ return AddNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const SubNode* op) {
+ return SubNode::make(Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MulNode* op) {
+ return AddNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b)));
+ }
+
+ PrimExpr VisitExpr_(const DivNode* op) {
+ return DivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const FloorDivNode* op) {
+ return FloorDivNode::make(
+ SubNode::make(
+ MulNode::make(Mutate(op->a), op->b),
+ MulNode::make(op->a, Mutate(op->b))),
+ MulNode::make(op->b, op->b));
+ }
+
+ PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const MinNode* op) {
+ return SelectNode::make(LENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const MaxNode* op) {
+ return SelectNode::make(GENode::make(op->a, op->b),
+ Mutate(op->a), Mutate(op->b));
+ }
+
+ PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED
+ PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED
+
+ PrimExpr VisitExpr_(const ReduceNode* op) {
+ // This case is relatively difficult because a reduction expression
Review comment:
seems to be a bit difficult. do you have an concrete example in mind?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services