This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 603f8bd721 [Unity][Op] Negative Log Likelihood Loss (#14517)
603f8bd721 is described below
commit 603f8bd721dc2756115aac5f113958cbe8ac2ce3
Author: Chaofan Lin <[email protected]>
AuthorDate: Fri Apr 7 00:10:13 2023 +0800
[Unity][Op] Negative Log Likelihood Loss (#14517)
This PR adds support for Negative Log Likelihood Loss
---
include/tvm/relax/attrs/nn.h | 13 +
python/tvm/relax/op/nn/nn.py | 45 +++
python/tvm/relax/transform/legalize_ops/nn.py | 29 ++
src/relax/op/nn/nn.cc | 225 +++++++++++
src/relax/op/nn/nn.h | 4 +
tests/python/relax/test_op_nn.py | 437 +++++++++++++++++++++
.../python/relax/test_transform_legalize_ops_nn.py | 257 ++++++++++++
tests/python/relax/test_tvmscript_parser_op_nn.py | 44 +++
8 files changed, 1054 insertions(+)
diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h
index bcfe3207bc..36e15909e2 100644
--- a/include/tvm/relax/attrs/nn.h
+++ b/include/tvm/relax/attrs/nn.h
@@ -285,6 +285,19 @@ struct GroupNormAttrs : public
tvm::AttrsNode<GroupNormAttrs> {
}
}; // struct GroupNormAttrs
+/*! \brief Attributes used in nll_loss operator */
+struct NLLLossAttrs : public tvm::AttrsNode<NLLLossAttrs> {
+ String reduction;
+ int ignore_index;
+
+ TVM_DECLARE_ATTRS(NLLLossAttrs, "relax.attrs.NLLLossAttrs") {
+ TVM_ATTR_FIELD(reduction).set_default("mean").describe(
+ "The reduction method to apply to the output. Can be"
+ "'none', 'mean' or 'sum'.");
+ TVM_ATTR_FIELD(ignore_index).describe("The target value to ignore.");
+ }
+}; // struct NLLLossAttrs
+
/*! \brief Attributes used in dropout operator */
struct DropoutAttrs : public tvm::AttrsNode<DropoutAttrs> {
double rate;
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index 02468637e0..cd1dfe1fb0 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -914,6 +914,51 @@ def cross_entropy_with_logits(predictions: Expr, labels:
Expr) -> Expr:
return _ffi_api.cross_entropy_with_logits(predictions, labels) # type:
ignore
+def nll_loss(
+ predictions: Expr,
+ targets: Expr,
+ weights: Optional[Expr] = None,
+ reduction: str = "mean",
+ ignore_index: int = -100,
+) -> Expr:
+ """Negative log likelihood loss.
+
+ `output[n, i_1, i_2, ..., i_k] = -p * w`, where
+ - `p = predictions[n, t, i_1, i_2, i_k]`,
+ - `t = targets[n, i_1, i_2, ..., i_k]`,
+ - `w = weights[t] if t != ignore_index else 0`
+
+ result = reduction(output)
+
+ Parameters
+ ----------
+ predictions : relax.Expr
+ The predictions. Should be a `(k+2)-D` Tensor with shape `(N, C, d_1,
d_2, ..., d_k)` where C
+ is the number of target classes.
+
+ targets : relax.Expr
+ The target value of each prediction. Should be a `(k+1)-D` Tensor with
shape
+ `(N, d_1, d_2, ..., d_k)`. Must be of int dtype.
+
+ weights : Optional[relax.Expr]
+ The weight of each target value. Should be a `1-D` Tensor with shape
`(C,)`.
+ If not specified, it is treated as if having all ones.
+
+ reduction : str
+ The reduction method to apply to the output.
+ Possible values are "mean", "sum" and "none".
+
+ ignore_index : int
+ The target value to ignore.
+
+ Returns
+ -------
+ result : relax.Expr
+ The computed result.
+ """
+ return _ffi_api.nll_loss(predictions, targets, weights, reduction,
ignore_index) # type: ignore
+
+
def attention(
query: Expr,
key: Expr,
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py
b/python/tvm/relax/transform/legalize_ops/nn.py
index 1ce4520635..59146186f9 100644
--- a/python/tvm/relax/transform/legalize_ops/nn.py
+++ b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -368,3 +368,32 @@ def _nn_attention_bias(bb: BlockBuilder, call: Call) ->
Expr:
call.attrs.scale,
primfunc_name_hint="attention_bias",
)
+
+
+@register_legalize("relax.nn.nll_loss")
+def _nn_nll_loss(bb: BlockBuilder, call: Call) -> Expr:
+ def nll_loss_without_weight(predictions, targets, reduction, ignore_index):
+ weight = topi.full(
+ (predictions.shape[1] if len(predictions.shape) > 1 else
predictions.shape[0],),
+ predictions.dtype,
+ 1.0,
+ )
+ return topi.nn.nll_loss(predictions, targets, weight, reduction,
ignore_index)
+
+ if len(call.args) == 2:
+ return bb.call_te(
+ nll_loss_without_weight,
+ call.args[0],
+ call.args[1],
+ reduction=call.attrs.reduction,
+ ignore_index=call.attrs.ignore_index,
+ )
+
+ return bb.call_te(
+ topi.nn.nll_loss,
+ call.args[0],
+ call.args[1],
+ call.args[2],
+ reduction=call.attrs.reduction,
+ ignore_index=call.attrs.ignore_index,
+ )
diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
index c3e18f8e3b..54faf3ee84 100644
--- a/src/relax/op/nn/nn.cc
+++ b/src/relax/op/nn/nn.cc
@@ -492,5 +492,230 @@ TVM_REGISTER_OP("relax.nn.cross_entropy_with_logits")
.add_argument("labels", "Tensor", "The labels.")
.set_attr<FInferStructInfo>("FInferStructInfo",
InferStructInfoCrossEntropy);
+/* relax.nn.nll_loss */
+TVM_REGISTER_NODE_TYPE(NLLLossAttrs);
+
+Expr nll_loss(Expr predictions, Expr targets, Optional<Expr> weights, String
reduction,
+ int ignore_index) {
+ ObjectPtr<NLLLossAttrs> attrs = make_object<NLLLossAttrs>();
+
+ ICHECK(reduction == "none" || reduction == "sum" || reduction == "mean")
+ << "The argument reduction of NLLLoss should be one of the following "
+ "values: none, mean, sum. However, the given value is "
+ << reduction;
+
+ attrs->reduction = std::move(reduction);
+ attrs->ignore_index = ignore_index;
+
+ static const Op& op = Op::Get("relax.nn.nll_loss");
+ if (weights.defined()) {
+ return Call(op, {std::move(predictions), std::move(targets),
std::move(weights.value())},
+ Attrs{attrs}, {});
+ } else {
+ return Call(op, {std::move(predictions), std::move(targets)},
Attrs{attrs}, {});
+ }
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.nll_loss").set_body_typed(nll_loss);
+
+StructInfo InferStructInfoNLLLoss(const Call& call, const BlockBuilder& ctx) {
+ if (call->args.size() < 2 || call->args.size() > 3) {
+ ctx->ReportFatal(Diagnostic::Error(call) << "NLLLoss op should take 2 or 3
arguments");
+ }
+
+ const auto* pred_sinfo =
GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
+ const auto* tgt_sinfo = GetStructInfoAs<TensorStructInfoNode>(call->args[1]);
+ const TensorStructInfoNode* wgt_sinfo = nullptr;
+ if (call->args.size() == 3) {
+ wgt_sinfo = GetStructInfoAs<TensorStructInfoNode>(call->args[2]);
+ if (wgt_sinfo == nullptr) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "NLLLoss requires the argument weights to be Tensor. However, the
given one is "
+ << call->args[2]->struct_info_->GetTypeKey());
+ }
+ }
+
+ if (pred_sinfo == nullptr) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "NLLLoss requires the argument preditions to be Tensor. However,
the given one is "
+ << call->args[0]->struct_info_->GetTypeKey());
+ }
+ if (tgt_sinfo == nullptr) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "NLLLoss requires the argument targets to be Tensor. However, the
given one is "
+ << call->args[1]->struct_info_->GetTypeKey());
+ }
+
+ // infer dtype
+ DataType output_dtype;
+ if (wgt_sinfo != nullptr) {
+ output_dtype = InferBinaryArithOpOutDtype(call, ctx,
GetRef<TensorStructInfo>(pred_sinfo),
+
GetRef<TensorStructInfo>(wgt_sinfo));
+ } else {
+ output_dtype = pred_sinfo->dtype;
+ }
+
+ // the type of targets must be int/uint.
+ if (!tgt_sinfo->IsUnknownDtype() && !tgt_sinfo->dtype.is_int() &&
!tgt_sinfo->dtype.is_uint()) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "NLLLoss expects the dtype of targets to be int/uint. However, the
dtype of targets is "
+ << tgt_sinfo->dtype);
+ }
+
+ // infer ndim
+ int K = kUnknownNDim; // k dim
+ if (!pred_sinfo->IsUnknownNdim()) {
+ if (pred_sinfo->ndim < 1) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "NLLLoss expects the ndim of predictions >= 1. However, the ndim
of predictions is "
+ << pred_sinfo->ndim);
+ }
+ K = pred_sinfo->ndim <= 2 ? 0 : pred_sinfo->ndim - 2;
+ }
+ if (!tgt_sinfo->IsUnknownNdim()) {
+ int K_tgt = tgt_sinfo->ndim <= 1 ? 0 : tgt_sinfo->ndim - 1;
+ if (K != kUnknownNDim && K != K_tgt) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "NLLLoss expects number of dimensions K inferred
from different "
+ "arguments to be equal. However, K from predictions
is "
+ << K << " while K from targets is " << K_tgt);
+ }
+ }
+ if (wgt_sinfo != nullptr && !wgt_sinfo->IsUnknownNdim() && wgt_sinfo->ndim
!= 1) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "NLLLoss expects the ndim of weights == 1. However,
the ndim of weights is "
+ << wgt_sinfo->ndim);
+ }
+
+ arith::Analyzer* analyzer = ctx->GetAnalyzer();
+ Optional<PrimExpr> N;
+ Optional<PrimExpr> C;
+ Array<PrimExpr> output_shape; // N, d1, d2, ..., dk
+
+ Optional<Array<PrimExpr>> pred_shape_value;
+ if (pred_sinfo->shape.defined()) {
+ pred_shape_value =
GetStructInfoAs<ShapeStructInfoNode>(pred_sinfo->shape.value())->values;
+ }
+ if (pred_shape_value.defined()) {
+ if (pred_shape_value.value().size() == 1) {
+ // (C,)
+ ICHECK(pred_sinfo->ndim == 1);
+ C = pred_shape_value.value()[0];
+ } else {
+ // (N, C, d1, d2, ..., dk)
+ ICHECK(pred_shape_value.value().size() >= 2);
+ ICHECK(pred_sinfo->ndim ==
static_cast<int>(pred_shape_value.value().size()));
+ N = pred_shape_value.value()[0];
+ C = pred_shape_value.value()[1];
+ output_shape = Array<PrimExpr>();
+ output_shape.push_back(N.value());
+ for (size_t i = 2; i < pred_shape_value.value().size(); ++i) {
+ output_shape.push_back(pred_shape_value.value()[i]);
+ }
+ }
+ }
+
+ Optional<Array<PrimExpr>> tgt_shape_value;
+ if (tgt_sinfo->shape.defined()) {
+ tgt_shape_value =
GetStructInfoAs<ShapeStructInfoNode>(tgt_sinfo->shape.value())->values;
+ }
+ if (tgt_shape_value.defined()) {
+ if (tgt_shape_value.value().empty()) {
+ // ()
+ ICHECK(tgt_sinfo->ndim == 0);
+ if (N.defined()) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "Shape mismatch for NLLLoss. Predictions shape is "
+ "(N, C, ...) while targets is a scalar");
+ }
+ } else {
+ // (N,) or (N, d1, d2, ..., dk)
+ // check N
+ const PrimExpr& N_tgt = tgt_shape_value.value()[0];
+ if (N.defined() && analyzer->CanProve(N.value() != N_tgt)) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "NLLLoss expects minibatch size N inferred from
different "
+ "arguments to be equal. However, N from
predictions is "
+ << N << " while N from targets is " << N_tgt);
+ }
+ // only C case
+ if (!N.defined() && C.defined()) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "Shape mismatch for NLLLoss. Predictions shape is "
+ "(C,) while targets is not a scalar");
+ }
+
+ if (tgt_shape_value.value().size() == 1) {
+ // (N,)
+ ICHECK(tgt_sinfo->IsUnknownNdim() || tgt_sinfo->ndim == 1);
+ } else {
+ // (N, d1, d2, ..., dk)
+ ICHECK(tgt_shape_value.value().size() >= 2);
+ ICHECK(tgt_sinfo->IsUnknownNdim() ||
+ tgt_sinfo->ndim ==
static_cast<int>(tgt_shape_value.value().size()));
+
+ if (pred_shape_value.defined()) {
+ // check (d1, d2, ..., dk)
+ for (size_t i = 1; i < tgt_shape_value.value().size(); ++i) {
+ if (analyzer->CanProve(output_shape[i] !=
tgt_shape_value.value()[i])) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "Shape mismatch for NLLLoss. The prediction
shape at this dim is "
+ << output_shape[i] << " while the target shape
at this dim is "
+ << tgt_shape_value.value()[i]);
+ }
+ }
+ }
+ }
+ }
+ }
+
+ if (wgt_sinfo != nullptr) {
+ Optional<Array<PrimExpr>> wgt_shape_value;
+ if (wgt_sinfo->shape.defined()) {
+ wgt_shape_value =
GetStructInfoAs<ShapeStructInfoNode>(wgt_sinfo->shape.value())->values;
+ }
+ if (wgt_shape_value.defined()) {
+ ICHECK(wgt_shape_value.value().size() == 1);
+ ICHECK(wgt_sinfo->IsUnknownNdim() || wgt_sinfo->ndim == 1);
+ const PrimExpr& C_wgt = wgt_shape_value.value()[0];
+ if (C.defined() && analyzer->CanProve(C.value() != C_wgt)) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "NLLLoss expects number of classes C inferred from
different "
+ "arguments to be equal. However, C from
predictions is "
+ << C << " while C from weights is " << C_wgt);
+ }
+ }
+ }
+
+ const auto* attrs = call->attrs.as<NLLLossAttrs>();
+ String reduction = attrs->reduction;
+
+ if (reduction == "none") {
+ // () or (N,) or (N, d1, d2, ..., dk)
+ if (pred_sinfo->shape.as<ShapeExprNode>()) {
+ return TensorStructInfo(ShapeExpr(output_shape), output_dtype);
+ } else {
+ int output_ndim = pred_sinfo->ndim == kUnknownNDim ? kUnknownNDim :
pred_sinfo->ndim - 1;
+ return TensorStructInfo(output_dtype, /*ndim=*/output_ndim);
+ }
+ } else {
+ // sum or mean. output is scalar
+ return TensorStructInfo(/*shape=*/ShapeExpr(Array<PrimExpr>()),
output_dtype);
+ }
+}
+
+TVM_REGISTER_OP("relax.nn.nll_loss")
+ .set_attrs_type<NLLLossAttrs>()
+ .set_num_inputs(3)
+ .add_argument("predictions", "Tensor", "The prediction tensor.")
+ .add_argument("targets", "Tensor", "The target tensor.")
+ .add_argument("weights", "Optional<Tensor>", "The weight of each target
values.")
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoNLLLoss);
+
} // namespace relax
} // namespace tvm
diff --git a/src/relax/op/nn/nn.h b/src/relax/op/nn/nn.h
index f578f89346..f0962f3018 100644
--- a/src/relax/op/nn/nn.h
+++ b/src/relax/op/nn/nn.h
@@ -85,6 +85,10 @@ Expr dropout(Expr data, double rate);
/*! \brief CrossEntropy with logits. */
Expr cross_entropy_with_logits(Expr predictions, Expr labels);
+/*! \brief Negative log likelihood loss. */
+Expr nll_loss(Expr predictions, Expr targets, Optional<Expr> weights, String
reduction,
+ int ignore_index);
+
} // namespace relax
} // namespace tvm
diff --git a/tests/python/relax/test_op_nn.py b/tests/python/relax/test_op_nn.py
index 5114478463..1b811ba0e5 100644
--- a/tests/python/relax/test_op_nn.py
+++ b/tests/python/relax/test_op_nn.py
@@ -1320,5 +1320,442 @@ def
test_cross_entropy_infer_struct_info_wrong_input_type():
bb.normalize(relax.op.nn.cross_entropy_with_logits(x1, y))
+def test_nll_loss_infer_struct_info():
+ bb = relax.BlockBuilder()
+
+ x0 = relax.Var("x", R.Tensor((3, 5, 10, 10), "float32"))
+ x1 = relax.Var("x", R.Tensor("float32", ndim=4))
+ x2 = relax.Var("x", R.Tensor("float32"))
+ x3 = relax.Var("x", R.Tensor((3, 5, 10, 10)))
+ x4 = relax.Var("x", R.Tensor((3, 5), "float32")) # (N, C)
+ x5 = relax.Var("x", R.Tensor((5,), "float32")) # (C,)
+
+ y0 = relax.Var("y", R.Tensor((3, 10, 10), "int64"))
+ y1 = relax.Var("y", R.Tensor("int64", ndim=3))
+ y2 = relax.Var("y", R.Tensor("int64"))
+ y3 = relax.Var("y", R.Tensor((3, 10, 10)))
+ y4 = relax.Var("y", R.Tensor((3,))) # (N,)
+ y5 = relax.Var("y", R.Tensor(())) # ()
+
+ w0 = relax.Var("w", R.Tensor((5,), "float32"))
+ w1 = relax.Var("w", R.Tensor("float32", ndim=1))
+ w2 = relax.Var("w", R.Tensor("float32"))
+ w3 = relax.Var("w", R.Tensor((5,)))
+
+ # reduction = mean
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x1, y0, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x2, y0, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x3, y0, w0, reduction="mean"),
+ relax.TensorStructInfo((), ""),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y1, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y2, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y3, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w1, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w2, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w3, reduction="mean"),
+ relax.TensorStructInfo((), ""),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x4, y4, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x5, y5, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+
+ # reduction=sum is totally the same as mean. Just need one test to ensure
they behave the same
+ _check_inference(
+ bb, relax.op.nn.nll_loss(x0, y0, w0, reduction="sum"),
relax.TensorStructInfo((), "float32")
+ )
+
+ # reduction=none
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w0, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x1, y0, w0, reduction="none"),
+ relax.TensorStructInfo(dtype="float32", ndim=3),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x2, y0, w0, reduction="none"),
+ relax.TensorStructInfo(dtype="float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x3, y0, w0, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), ""),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y1, w0, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y2, w0, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y3, w0, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w1, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w2, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w3, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), ""),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x4, y4, w0, reduction="none"),
+ relax.TensorStructInfo((3,), "float32"), # (N,)
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x5, y5, w0, reduction="none"),
+ relax.TensorStructInfo((), "float32"), # ()
+ )
+
+
+def test_nll_loss_infer_struct_info_shape_symbolic():
+ bb = relax.BlockBuilder()
+ N = tir.Var("N", "int64")
+ C = tir.Var("C", "int64")
+ d1 = tir.Var("d", "int64")
+ d2 = tir.Var("d", "int64")
+ x0 = relax.Var("x", R.Tensor((N, C, d1, d2), "float32"))
+ x1 = relax.Var("x", R.Tensor((N, C), "float32"))
+ x2 = relax.Var("x", R.Tensor((C,), "float32"))
+ x3 = relax.Var("x", R.Tensor((3, C, d1, 2), "float32"))
+ y0 = relax.Var("y", R.Tensor((N, d1, d2), "int64"))
+ y1 = relax.Var("y", R.Tensor((N,), "int64"))
+ y2 = relax.Var("y", R.Tensor((), "int64"))
+ y3 = relax.Var("y", R.Tensor((3, d1, 2), "int64"))
+ w0 = relax.Var("w", R.Tensor((C,), "float32"))
+ w1 = relax.Var("w", R.Tensor((5,), "float32"))
+
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w0, reduction="none"),
+ relax.TensorStructInfo((N, d1, d2), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x1, y1, w0, reduction="none"),
+ relax.TensorStructInfo((N,), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x2, y2, w0, reduction="none"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x3, y3, w0, reduction="none"),
+ relax.TensorStructInfo((3, d1, 2), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x3, y3, w1, reduction="none"),
+ relax.TensorStructInfo((3, d1, 2), "float32"),
+ )
+
+
+def test_nll_loss_infer_struct_info_shape_var():
+ bb = relax.BlockBuilder()
+
+ s0 = relax.Var("s0", relax.ShapeStructInfo((3, 5, 10, 10)))
+ s1 = relax.Var("s1", relax.ShapeStructInfo(ndim=4))
+ s2 = relax.Var("s2", relax.ShapeStructInfo())
+ s3 = relax.Var("s3", relax.ShapeStructInfo((3, 10, 10)))
+ s4 = relax.Var("s4", relax.ShapeStructInfo(ndim=3))
+ s5 = relax.Var("s5", relax.ShapeStructInfo((5,)))
+ s6 = relax.Var("s6", relax.ShapeStructInfo(ndim=1))
+
+ x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+ x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+ x2 = relax.Var("x", relax.TensorStructInfo(s2, "float32"))
+ y0 = relax.Var("y", relax.TensorStructInfo(s3, "int64"))
+ y1 = relax.Var("y", relax.TensorStructInfo(s4, "int64"))
+ w0 = relax.Var("w", relax.TensorStructInfo(s5, "float32"))
+ w1 = relax.Var("w", relax.TensorStructInfo(s6, "float32"))
+
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w0, reduction="none"),
+ relax.TensorStructInfo(dtype="float32", ndim=3),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x1, y0, w0, reduction="none"),
+ relax.TensorStructInfo(dtype="float32", ndim=3),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x2, y0, w0, reduction="none"),
+ relax.TensorStructInfo(dtype="float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y1, w0, reduction="none"),
+ relax.TensorStructInfo(dtype="float32", ndim=3),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w1, reduction="none"),
+ relax.TensorStructInfo(dtype="float32", ndim=3),
+ )
+
+
+def test_nll_loss_infer_struct_info_no_weights():
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", R.Tensor((3, 5, 10, 10), "float32"))
+ y = relax.Var("x", R.Tensor((3, 10, 10), "int64"))
+
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x, y, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x, y, reduction="none"),
+ relax.TensorStructInfo((3, 10, 10), "float32"),
+ )
+
+
+def test_nll_loss_infer_struct_info_no_weights_symbolic():
+ N = tir.Var("N", "int64")
+ C = tir.Var("C", "int64")
+ d1 = tir.Var("d", "int64")
+ d2 = tir.Var("d", "int64")
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", R.Tensor((N, C, d1, d2), "float32"))
+ y = relax.Var("y", R.Tensor((N, d1, d2), "int64"))
+
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x, y, reduction="mean"),
+ relax.TensorStructInfo((), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x, y, reduction="none"),
+ relax.TensorStructInfo((N, d1, d2), "float32"),
+ )
+
+
+def test_nll_loss_infer_struct_info_wrong_input_type():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((3, 5, 10, 10), "float32"))
+ x1 = relax.Var("x", relax.ShapeStructInfo((2, 3)))
+ x2 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3), "float32")))
+ y0 = relax.Var("y", R.Tensor((3, 10, 10), "int64"))
+ y1 = relax.Var("y", relax.ShapeStructInfo((2, 3)))
+ y2 = relax.Var("y", relax.FuncStructInfo([], R.Tensor((2, 3), "float32")))
+ w0 = relax.Var("w", R.Tensor((5,), "float32"))
+ w1 = relax.Var("w", relax.ShapeStructInfo((2, 3)))
+ w2 = relax.Var("w", relax.FuncStructInfo([], R.Tensor((2, 3), "float32")))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x1, y0, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x2, y0, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y1, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y2, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y0, w1))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y0, w2))
+
+
+def test_nll_loss_infer_struct_info_more_input_dtype():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((3, 5, 10, 10), "float16"))
+ x1 = relax.Var("x", R.Tensor((3, 5, 10, 10), "int8"))
+ x2 = relax.Var("x", R.Tensor((3, 5, 10, 10), "int32"))
+ x3 = relax.Var("x", R.Tensor((3, 5, 10, 10), "float64"))
+ y0 = relax.Var("y", R.Tensor((3, 10, 10), "int8"))
+ w0 = relax.Var("y", R.Tensor((5,), "float16"))
+ w1 = relax.Var("y", R.Tensor((5,), "int8"))
+ w2 = relax.Var("y", R.Tensor((5,), "int32"))
+ w3 = relax.Var("y", R.Tensor((5,), "float64"))
+
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x0, y0, w0, reduction="mean"),
+ relax.TensorStructInfo((), "float16"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x1, y0, w1, reduction="mean"),
+ relax.TensorStructInfo((), "int8"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x2, y0, w2, reduction="mean"),
+ relax.TensorStructInfo((), "int32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.nn.nll_loss(x3, y0, w3, reduction="mean"),
+ relax.TensorStructInfo((), "float64"),
+ )
+
+
+def test_nll_loss_infer_struct_info_targets_dtype():
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", R.Tensor((3, 5, 10, 10), "float32"))
+ w = relax.Var("w", R.Tensor((5,), "float32"))
+ targets0 = relax.Var("targets", R.Tensor((3, 10, 10), "float32"))
+ targets1 = relax.Var("targets", R.Tensor((3, 10, 10), "float64"))
+ targets2 = relax.Var("targets", R.Tensor((3, 10, 10), "bool"))
+ targets3 = relax.Var("targets", R.Tensor((3, 10, 10), "int32"))
+ targets4 = relax.Var("targets", R.Tensor((3, 10, 10), "int64"))
+ targets5 = relax.Var("targets", R.Tensor((3, 10, 10), "uint32"))
+ targets6 = relax.Var("targets", R.Tensor((3, 10, 10), ""))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x, targets0, w))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x, targets1, w))
+
+ # correct cases
+ bb.normalize(relax.op.nn.nll_loss(x, targets2, w)) # bool is uint1
+ bb.normalize(relax.op.nn.nll_loss(x, targets3, w))
+ bb.normalize(relax.op.nn.nll_loss(x, targets4, w))
+ bb.normalize(relax.op.nn.nll_loss(x, targets5, w))
+ bb.normalize(relax.op.nn.nll_loss(x, targets6, w)) # unknwon dtype
+
+
+def test_nll_loss_infer_struct_info_ndim_mismatch():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((3, 5, 10, 10), "float32"))
+ x1 = relax.Var("x", R.Tensor((3, 5, 10, 10, 10), "float32"))
+ x2 = relax.Var("x", R.Tensor((3, 5, 10), "float32"))
+ y0 = relax.Var("x", R.Tensor((3, 10, 10), "int64"))
+ y1 = relax.Var("x", R.Tensor((3, 10, 10, 10), "int64"))
+ y2 = relax.Var("x", R.Tensor((3, 10), "int64"))
+ w0 = relax.Var("w", R.Tensor((5,), "float32"))
+ w1 = relax.Var("w", R.Tensor((5, 5), "float32"))
+ w2 = relax.Var("w", R.Tensor((), "float32"))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x1, y0, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x2, y0, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y1, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y2, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y0, w1))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y0, w2))
+
+
+def test_nll_loss_infer_struct_info_shape_mismatch():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((3, 5, 10, 10), "float32"))
+ x1 = relax.Var("x", R.Tensor((3, 6, 10, 10), "float32"))
+ x2 = relax.Var("x", R.Tensor((4, 5, 10, 10), "float32"))
+ x3 = relax.Var("x", R.Tensor((3, 5, 11, 10), "float32"))
+ y0 = relax.Var("x", R.Tensor((3, 10, 10), "int64"))
+ y1 = relax.Var("x", R.Tensor((4, 10, 10), "int64"))
+ y2 = relax.Var("x", R.Tensor((3, 11, 10), "int64"))
+ w0 = relax.Var("w", R.Tensor((5,), "float32"))
+ w1 = relax.Var("w", R.Tensor((4,), "float32"))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x1, y0, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x2, y0, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x3, y0, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y1, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y2, w0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x0, y0, w1))
+
+
+def test_nll_loss_infer_struct_info_wrong_reduction():
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", R.Tensor((3, 5, 10, 10), "float32"))
+ y = relax.Var("x", R.Tensor((3, 10, 10), "int64"))
+ w = relax.Var("w", R.Tensor((5,), "float32"))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.nn.nll_loss(x, y, w, reduction="foo"))
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py
b/tests/python/relax/test_transform_legalize_ops_nn.py
index e807082e35..fe663527c3 100644
--- a/tests/python/relax/test_transform_legalize_ops_nn.py
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -2443,5 +2443,262 @@ def test_attention():
tvm.ir.assert_structural_equal(mod, Expected)
+def test_nll_loss():
+ # fmt: off
+ @tvm.script.ir_module
+ class NLLLoss:
+ @R.function
+ def main(predictions: R.Tensor((2, 3, 4, 5), "float32"), targets:
R.Tensor((2, 4, 5), "int64"), weights: R.Tensor((4,), "float32")) ->
R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.nn.nll_loss(predictions, targets,
weights, reduction="mean", ignore_index=-1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(predictions: R.Tensor((2, 3, 4, 5), dtype="float32"),
targets: R.Tensor((2, 4, 5), dtype="int64"), weights: R.Tensor((4,),
dtype="float32"),) -> R.Tensor((), dtype="float32"):
+ # block 0
+ gv = R.call_tir(Expected.nll_loss, (predictions, targets,
weights), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def nll_loss(rxplaceholder: T.Buffer((T.int64(2), T.int64(3),
T.int64(4), T.int64(5)), "float32"), rxplaceholder_1: T.Buffer((T.int64(2),
T.int64(4), T.int64(5)), "int64"), rxplaceholder_2: T.Buffer(T.int64(4),
"float32"), T_divide: T.Buffer((), "float32"),):
+ # function attr dict
+ T.func_attr({"tir.noalias": True})
+ # body
+ # with T.block("root")
+ nll_loss = T.alloc_buffer([T.int64(2), T.int64(4), T.int64(5)],
dtype="float32")
+ nll_loss_red = T.alloc_buffer([], dtype="float32")
+ nll_loss_1 = T.alloc_buffer([T.int64(2), T.int64(4), T.int64(5)],
dtype="float32")
+ nll_loss_red_1 = T.alloc_buffer([], dtype="float32")
+ for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(rxplaceholder_1[v_ax0, v_ax1, v_ax2],
rxplaceholder[v_ax0, rxplaceholder_1[v_ax0, v_ax1, v_ax2], v_ax1, v_ax2],
rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]])
+ T.writes(nll_loss[v_ax0, v_ax1, v_ax2])
+ nll_loss[v_ax0, v_ax1, v_ax2] =
T.Select(rxplaceholder_1[v_ax0, v_ax1, v_ax2] != T.int64(-1), (T.float32(0) -
rxplaceholder[v_ax0, rxplaceholder_1[v_ax0, v_ax1, v_ax2], v_ax1, v_ax2]) *
rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]], T.float32(0))
+ for k0, k1, k2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss_red"):
+ v_k0, v_k1, v_k2 = T.axis.remap("RRR", [k0, k1, k2])
+ T.reads(nll_loss[v_k0, v_k1, v_k2])
+ T.writes(nll_loss_red[()])
+ with T.init():
+ nll_loss_red[()] = T.float32(0)
+ nll_loss_red[()] = nll_loss_red[()] + nll_loss[v_k0, v_k1,
v_k2]
+ for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(rxplaceholder_1[v_ax0, v_ax1, v_ax2],
rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]])
+ T.writes(nll_loss_1[v_ax0, v_ax1, v_ax2])
+ nll_loss_1[v_ax0, v_ax1, v_ax2] =
T.Select(rxplaceholder_1[v_ax0, v_ax1, v_ax2] != T.int64(-1),
rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]], T.float32(0))
+ for k0, k1, k2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss_red_1"):
+ v_k0, v_k1, v_k2 = T.axis.remap("RRR", [k0, k1, k2])
+ T.reads(nll_loss_1[v_k0, v_k1, v_k2])
+ T.writes(nll_loss_red_1[()])
+ with T.init():
+ nll_loss_red_1[()] = T.float32(0)
+ nll_loss_red_1[()] = nll_loss_red_1[()] + nll_loss_1[v_k0,
v_k1, v_k2]
+ with T.block("T_divide"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(nll_loss_red[()], nll_loss_red_1[()])
+ T.writes(T_divide[()])
+ T_divide[()] = nll_loss_red[()] / nll_loss_red_1[()]
+ # fmt: on
+ mod = LegalizeOps()(NLLLoss)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_nll_no_weight():
+ # fmt: off
+ @tvm.script.ir_module
+ class NLLLoss:
+ @R.function
+ def main(predictions: R.Tensor((2, 3, 4, 5), "float32"), targets:
R.Tensor((2, 4, 5), "int64")) -> R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.nn.nll_loss(predictions, targets,
reduction="mean", ignore_index=-1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(predictions: R.Tensor((2, 3, 4, 5), dtype="float32"),
targets: R.Tensor((2, 4, 5), dtype="int64"),) -> R.Tensor((), dtype="float32"):
+ # block 0
+ gv = R.call_tir(Expected.nll_loss_without_weight, (predictions,
targets), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def nll_loss_without_weight(rxplaceholder: T.Buffer((T.int64(2),
T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_1:
T.Buffer((T.int64(2), T.int64(4), T.int64(5)), "int64"), T_divide: T.Buffer((),
"float32"),):
+ # function attr dict
+ T.func_attr({"tir.noalias": True})
+ # body
+ # with T.block("root")
+ T_full = T.alloc_buffer([T.int64(3)], dtype="float32")
+ nll_loss = T.alloc_buffer([T.int64(2), T.int64(4), T.int64(5)],
dtype="float32")
+ nll_loss_red = T.alloc_buffer([], dtype="float32")
+ nll_loss_1 = T.alloc_buffer([T.int64(2), T.int64(4), T.int64(5)],
dtype="float32")
+ nll_loss_red_1 = T.alloc_buffer([], dtype="float32")
+ for ax0 in T.serial(T.int64(3)):
+ with T.block("T_full"):
+ v_ax0 = T.axis.spatial(T.int64(3), ax0)
+ T.reads()
+ T.writes(T_full[v_ax0])
+ T_full[v_ax0] = T.float32(1)
+ for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(rxplaceholder_1[v_ax0, v_ax1, v_ax2],
rxplaceholder[v_ax0, rxplaceholder_1[v_ax0, v_ax1, v_ax2], v_ax1, v_ax2],
T_full[rxplaceholder_1[v_ax0, v_ax1, v_ax2]])
+ T.writes(nll_loss[v_ax0, v_ax1, v_ax2])
+ nll_loss[v_ax0, v_ax1, v_ax2] =
T.Select(rxplaceholder_1[v_ax0, v_ax1, v_ax2] != T.int64(-1), (T.float32(0) -
rxplaceholder[v_ax0, rxplaceholder_1[v_ax0, v_ax1, v_ax2], v_ax1, v_ax2]) *
T_full[rxplaceholder_1[v_ax0, v_ax1, v_ax2]], T.float32(0))
+ for k0, k1, k2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss_red"):
+ v_k0, v_k1, v_k2 = T.axis.remap("RRR", [k0, k1, k2])
+ T.reads(nll_loss[v_k0, v_k1, v_k2])
+ T.writes(nll_loss_red[()])
+ with T.init():
+ nll_loss_red[()] = T.float32(0)
+ nll_loss_red[()] = nll_loss_red[()] + nll_loss[v_k0, v_k1,
v_k2]
+ for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(rxplaceholder_1[v_ax0, v_ax1, v_ax2],
T_full[rxplaceholder_1[v_ax0, v_ax1, v_ax2]])
+ T.writes(nll_loss_1[v_ax0, v_ax1, v_ax2])
+ nll_loss_1[v_ax0, v_ax1, v_ax2] =
T.Select(rxplaceholder_1[v_ax0, v_ax1, v_ax2] != T.int64(-1),
T_full[rxplaceholder_1[v_ax0, v_ax1, v_ax2]], T.float32(0))
+ for k0, k1, k2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
+ with T.block("nll_loss_red_1"):
+ v_k0, v_k1, v_k2 = T.axis.remap("RRR", [k0, k1, k2])
+ T.reads(nll_loss_1[v_k0, v_k1, v_k2])
+ T.writes(nll_loss_red_1[()])
+ with T.init():
+ nll_loss_red_1[()] = T.float32(0)
+ nll_loss_red_1[()] = nll_loss_red_1[()] + nll_loss_1[v_k0,
v_k1, v_k2]
+ with T.block("T_divide"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(nll_loss_red[()], nll_loss_red_1[()])
+ T.writes(T_divide[()])
+ T_divide[()] = nll_loss_red[()] / nll_loss_red_1[()]
+ # fmt: on
+
+ mod = LegalizeOps()(NLLLoss)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_nll_no_batch():
+ # fmt: off
+ @tvm.script.ir_module
+ class NLLLoss:
+ @R.function
+ def main(predictions: R.Tensor(("C",), "float32"), targets:
R.Tensor((), "int64"), weights: R.Tensor(("C",), "float32")) -> R.Tensor((),
"float32"):
+ gv = R.nn.nll_loss(predictions, targets, weights,
reduction="mean", ignore_index=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(predictions: R.Tensor(("C",), dtype="float32"), targets:
R.Tensor((), dtype="int64"), weights: R.Tensor(("C",), dtype="float32")) ->
R.Tensor((), dtype="float32"):
+ C = T.int64()
+ gv = R.call_tir(Expected.nll_loss, (predictions, targets,
weights), out_sinfo=R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def nll_loss(var_rxplaceholder: T.handle, rxplaceholder: T.Buffer((),
"int64"), var_rxplaceholder_1: T.handle, T_divide: T.Buffer((), "float32")):
+ T.func_attr({"tir.noalias": True})
+ C = T.int64()
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder, (C,))
+ rxplaceholder_2 = T.match_buffer(var_rxplaceholder_1, (C,))
+ # with T.block("root"):
+ nll_loss = T.alloc_buffer(())
+ nll_loss_1 = T.alloc_buffer(())
+ with T.block("nll_loss"):
+ vi = T.axis.spatial(T.int64(1), T.int64(0))
+ T.reads(rxplaceholder[()], rxplaceholder_1[rxplaceholder[()]],
rxplaceholder_2[rxplaceholder[()]])
+ T.writes(nll_loss[()])
+ nll_loss[()] = T.Select(rxplaceholder[()] != T.int64(1),
(T.float32(0) - rxplaceholder_1[rxplaceholder[()]]) *
rxplaceholder_2[rxplaceholder[()]], T.float32(0))
+ with T.block("nll_loss_1"):
+ vi = T.axis.spatial(T.int64(1), T.int64(0))
+ T.reads(rxplaceholder[()], rxplaceholder_2[rxplaceholder[()]])
+ T.writes(nll_loss_1[()])
+ nll_loss_1[()] = T.Select(rxplaceholder[()] != T.int64(1),
rxplaceholder_2[rxplaceholder[()]], T.float32(0))
+ with T.block("T_divide"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(nll_loss[()], nll_loss_1[()])
+ T.writes(T_divide[()])
+ T_divide[()] = nll_loss[()] / nll_loss_1[()]
+ # fmt: on
+
+ mod = LegalizeOps()(NLLLoss)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_nll_loss_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class NLLLoss:
+ @R.function
+ def main(predictions: R.Tensor(("N", "C", "d1", "d2"), "float32"),
targets: R.Tensor(("N", "d1", "d2"), "int64"), weights: R.Tensor(("C",),
"float32")) -> R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.nn.nll_loss(predictions, targets,
weights, reduction="mean", ignore_index=-1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(predictions: R.Tensor(("N", "C", "d1", "d2"),
dtype="float32"), targets: R.Tensor(("N", "d1", "d2"), dtype="int64"), weights:
R.Tensor(("C",), dtype="float32")) -> R.Tensor((), dtype="float32"):
+ # block 0
+ gv = R.call_tir(Expected.nll_loss, (predictions, targets,
weights), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def nll_loss(var_rxplaceholder: T.handle, var_rxplaceholder_1:
T.handle, var_rxplaceholder_2: T.handle, T_divide: T.Buffer((), "float32"),):
+ # function attr dict
+ T.func_attr({"tir.noalias": True})
+ C = T.int64()
+ N = T.int64()
+ d1 = T.int64()
+ d2 = T.int64()
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [N, C, d1, d2],
dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [N, d1, d2],
dtype="int64")
+ rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [C],
dtype="float32")
+ # body
+ # with T.block("root")
+ nll_loss = T.alloc_buffer([N, d1, d2], dtype="float32")
+ nll_loss_red = T.alloc_buffer([], dtype="float32")
+ nll_loss_1 = T.alloc_buffer([N, d1, d2], dtype="float32")
+ nll_loss_red_1 = T.alloc_buffer([], dtype="float32")
+ for ax0, ax1, ax2 in T.grid(N, d1, d2):
+ with T.block("nll_loss"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(rxplaceholder_1[v_ax0, v_ax1, v_ax2],
rxplaceholder[v_ax0, rxplaceholder_1[v_ax0, v_ax1, v_ax2], v_ax1,
v_ax2],rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]],)
+ T.writes(nll_loss[v_ax0, v_ax1, v_ax2])
+ nll_loss[v_ax0, v_ax1, v_ax2] =
T.Select(rxplaceholder_1[v_ax0, v_ax1, v_ax2] != T.int64(-1), (T.float32(0) -
rxplaceholder[v_ax0, rxplaceholder_1[v_ax0, v_ax1, v_ax2], v_ax1, v_ax2]) *
rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]], T.float32(0),)
+ for k0, k1, k2 in T.grid(N, d1, d2):
+ with T.block("nll_loss_red"):
+ v_k0, v_k1, v_k2 = T.axis.remap("RRR", [k0, k1, k2])
+ T.reads(nll_loss[v_k0, v_k1, v_k2])
+ T.writes(nll_loss_red[()])
+ with T.init():
+ nll_loss_red[()] = T.float32(0)
+ nll_loss_red[()] = nll_loss_red[()] + nll_loss[v_k0, v_k1,
v_k2]
+ for ax0, ax1, ax2 in T.grid(N, d1, d2):
+ with T.block("nll_loss_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(rxplaceholder_1[v_ax0, v_ax1, v_ax2],
rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]],)
+ T.writes(nll_loss_1[v_ax0, v_ax1, v_ax2])
+ nll_loss_1[v_ax0, v_ax1, v_ax2] =
T.Select(rxplaceholder_1[v_ax0, v_ax1, v_ax2] != T.int64(-1),
rxplaceholder_2[rxplaceholder_1[v_ax0, v_ax1, v_ax2]], T.float32(0),)
+ for k0, k1, k2 in T.grid(N, d1, d2):
+ with T.block("nll_loss_red_1"):
+ v_k0, v_k1, v_k2 = T.axis.remap("RRR", [k0, k1, k2])
+ T.reads(nll_loss_1[v_k0, v_k1, v_k2])
+ T.writes(nll_loss_red_1[()])
+ with T.init():
+ nll_loss_red_1[()] = T.float32(0)
+ nll_loss_red_1[()] = nll_loss_red_1[()] + nll_loss_1[v_k0,
v_k1, v_k2]
+ with T.block("T_divide"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(nll_loss_red[()], nll_loss_red_1[()])
+ T.writes(T_divide[()])
+ T_divide[()] = nll_loss_red[()] / nll_loss_red_1[()]
+ # fmt: on
+ mod = LegalizeOps()(NLLLoss)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_tvmscript_parser_op_nn.py
b/tests/python/relax/test_tvmscript_parser_op_nn.py
index a822fae719..014d524751 100644
--- a/tests/python/relax/test_tvmscript_parser_op_nn.py
+++ b/tests/python/relax/test_tvmscript_parser_op_nn.py
@@ -302,5 +302,49 @@ def test_cross_entropy_with_logits():
_check(foo, bb.get()["foo"])
+def test_nll_loss():
+ @R.function
+ def foo(
+ predictions: R.Tensor((3, 5, 10, 10), dtype="float32"),
+ targets: R.Tensor((3, 10, 10), dtype="int64"),
+ weights: R.Tensor((5,), dtype="float32"),
+ ) -> R.Tensor((), dtype="float32"):
+ gv: R.Tensor((), dtype="float32") = R.nn.nll_loss(predictions,
targets, weights, "mean", -1)
+ return gv
+
+ predictions = relax.Var("predictions", R.Tensor((3, 5, 10, 10), "float32"))
+ targets = relax.Var("targets", R.Tensor((3, 10, 10), "int64"))
+ weights = relax.Var("weights", R.Tensor((5,), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", [predictions, targets, weights]):
+ gv = bb.emit(
+ relax.op.nn.nll_loss(predictions, targets, weights,
reduction="mean", ignore_index=-1)
+ )
+ bb.emit_func_output(gv)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_nll_loss_no_weights():
+ @R.function
+ def foo(
+ predictions: R.Tensor((3, 5, 10, 10), dtype="float32"),
+ targets: R.Tensor((3, 10, 10), dtype="int64"),
+ ) -> R.Tensor((), dtype="float32"):
+ gv: R.Tensor((), dtype="float32") = R.nn.nll_loss(
+ predictions, targets, reduction="mean", ignore_index=-1
+ )
+ return gv
+
+ predictions = relax.Var("predictions", R.Tensor((3, 5, 10, 10), "float32"))
+ targets = relax.Var("targets", R.Tensor((3, 10, 10), "int64"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", [predictions, targets]):
+ gv = bb.emit(relax.op.nn.nll_loss(predictions, targets,
reduction="mean", ignore_index=-1))
+ bb.emit_func_output(gv)
+
+ _check(foo, bb.get()["foo"])
+
+
if __name__ == "__main__":
tvm.testing.main()
