Searched refs:tensor_argmax (Results 1 – 1 of 1) sorted by relevance
126 Tensor<DenseIndex, 0, DataLayout> tensor_argmax; in test_simple_argmax() local128 tensor_argmax = tensor.argmax(); in test_simple_argmax()130 VERIFY_IS_EQUAL(tensor_argmax(0), 0); in test_simple_argmax()134 tensor_argmax = tensor.argmax(); in test_simple_argmax()136 VERIFY_IS_EQUAL(tensor_argmax(0), 2*3*5*7 - 1); in test_simple_argmax()170 Tensor<DenseIndex, 3, DataLayout> tensor_argmax; in test_argmax_dim() local185 tensor_argmax = tensor.argmax(dim); in test_argmax_dim()187 VERIFY_IS_EQUAL(tensor_argmax.size(), in test_argmax_dim()189 for (ptrdiff_t n = 0; n < tensor_argmax.size(); ++n) { in test_argmax_dim()191 VERIFY_IS_EQUAL(tensor_argmax.data()[n], 0); in test_argmax_dim()[all …]