| /external/tensorflow/tensorflow/core/kernels/ |
| argmax_op_gpu.cu.cc | 28 template struct functor::ArgMax<GPUDevice, T, int64>; \
30 template struct functor::ArgMax<GPUDevice, T, int32>; \
|
| argmax_op.cc | 107 : public ArgOp<Device, T, Tout, functor::ArgMax<Device, T, Tout> > {
110 : ArgOp<Device, T, Tout, functor::ArgMax<Device, T, Tout> >(context) {}
122 REGISTER_KERNEL_BUILDER(Name("ArgMax") \
134 REGISTER_KERNEL_BUILDER(Name("ArgMax") \
156 void ArgMax<GPUDevice, T, Tout>::Reduce##Dims( \
177 extern template struct ArgMax<GPUDevice, T, int64>; \
179 extern template struct ArgMax<GPUDevice, T, int32>; \
192 REGISTER_KERNEL_BUILDER(Name("ArgMax") \
206 REGISTER_KERNEL_BUILDER(Name("ArgMax") \
|
| argmax_op.h | 29 struct ArgMax {
34 output.device(d) = input.argmax(dimension).template cast<Tout>(); \
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| categorical_op.cc | 81 xla::ComputationDataHandle argmax; variable
84 XlaHelpers::ArgMax(builder, ctx, softmax_entries, softmax_shape,
85 input_type(0), output_type(0), /*axis=*/2, &argmax));
87 ctx->SetOutput(0, argmax);
|
| index_ops.cc | 74 XlaHelpers::ArgMax(b, ctx, input, input_shape, input_type(0),
83 REGISTER_XLA_OP(Name("ArgMax")
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| xla_helpers.h | 75 // Sets `argmax` to the argmax of `input` along `axis`. `input_shape` and
77 // `output_type` is the dtype to use for `argmax`.
78 static Status ArgMax(xla::ComputationBuilder* builder,
83 xla::ComputationDataHandle* argmax);
|
| xla_helpers.cc | 182 Status XlaHelpers::ArgMax(xla::ComputationBuilder* builder,
187 xla::ComputationDataHandle* argmax) {
189 axis, /*is_min=*/false, argmax);
|
| /external/tensorflow/tensorflow/core/common_runtime/ |
| shape_refiner_test.cc | 213 auto am = ops::ArgMax(root, input, dim);
230 auto am = ops::ArgMax(root, input, dim);
246 auto am = ops::ArgMax(root, input, dim);
[all...] |
| /external/tensorflow/tensorflow/contrib/lite/toco/ |
| tooling_util.cc | 315 HANDLE_OPERATORTYPENAME_CASE(ArgMax)
[all...] |
| /external/tensorflow/tensorflow/compiler/tests/ |
| randomized_tests.cc | [all...] |
| /external/tensorflow/tensorflow/contrib/lite/kernels/internal/reference/ |
| reference_ops.h | [all...] |
| /external/tensorflow/tensorflow/contrib/lite/kernels/internal/optimized/ |
| optimized_ops.h | [all...] |
