return lst
)JIT";
+const auto nonzero_tensor = R"JIT(
+ def forward(self, input: Tensor):
+ a = torch.nonzero(input).clone()
+ return (a)
+)JIT";
+
const std::string quantize_script = R"IR(
graph(%input: Tensor, %weights: Tensor):
%scale: float = prim::Constant[value=1.]()
testStaticRuntime(cat_script, args0, args1);
}
-
TEST(StaticRuntime, IndividualOps_Cumsum) {
auto a = at::randn({2, 3});
std::vector<IValue> args0{a, 0};
std::vector<IValue> args1{b, 1, dtype};
testStaticRuntime(cumsum_script_dtype, args0, args1);
}
+
+TEST(StaticRuntime, IndividualOps_Nonzero) {
+ auto a = at::randint(0, 2, {2, 3});
+ testStaticRuntime(nonzero_tensor, {a});
+
+ auto b = at::randint(0, 2, {4, 3, 2});
+ testStaticRuntime(nonzero_tensor, {a}, {b});
+}
};
});
+REGISTER_OPERATOR_FUNCTOR(
+ aten::nonzero,
+ aten_nonzero,
+ [](Node* n) -> SROperator {
+ if (!n->matches(torch::schema("aten::nonzero(Tensor self) -> Tensor"))) {
+ LogAndDumpSchema(n);
+ return nullptr;
+ }
+ return [](ProcessedNode* p_node) {
+ const auto& input = p_node->Input(0).toTensor();
+ if (p_node->Output(0).isNone()) {
+ p_node->Output(0) = at::native::nonzero_cpu(input);
+ return;
+ }
+
+ auto& output = p_node->Output(0).toTensor();
+ fastResizeToZero(output);
+ at::native::nonzero_out_cpu(input, output);
+ };
+ });
+
namespace {
void check_cat_no_zero_dim(const std::vector<at::Tensor>& tensors) {