Home
last modified time | relevance | path

Searched refs:k_tensor (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/
Din_topk_op.h39 const Tensor* k_tensor = nullptr; member
63 if (k.k_tensor != nullptr) {
64 if (k.k_tensor->dtype() == DT_INT32) {
65 k_val = k.k_tensor->scalar<int32>()();
67 k_val = k.k_tensor->scalar<int64_t>()();
Din_topk_op.cc45 const Tensor* k_tensor = nullptr; in Compute() local
54 k_tensor = &k_in; in Compute()
79 arg.k_tensor = k_tensor; in Compute()
Din_topk_op_gpu.cu.cc151 if (k.k_tensor != nullptr) { in operator ()()
152 if (k.k_tensor->dtype() == DT_INT32) { in operator ()()
155 (cnt < k.k_tensor->flat<int32>().template cast<int64_t>().broadcast( in operator ()()
160 (cnt < k.k_tensor->flat<int64_t>().broadcast(Dims<1>(num_targets))); in operator ()()
Dclustering_ops.cc293 const Tensor& k_tensor = context->input(2); in Compute() local
299 OP_REQUIRES(context, TensorShapeUtils::IsScalar(k_tensor.shape()), in Compute()
319 std::min<int64_t>(num_centers, k_tensor.scalar<int64_t>()()); in Compute()
/external/tensorflow/tensorflow/python/compiler/tensorrt/test/
Dtopk_test.py30 k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
31 values, indices = nn_ops.top_k(x, k_tensor, name="TopK")
51 k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
52 values, indices = nn_ops.top_k(x, k_tensor, name="TopK")