Searched refs:axis_type (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/lite/kernels/ |
| D | arg_min_max_test.cc | 33 ArgBaseOpModel(TensorType input_type, int axis_value, TensorType axis_type, in ArgBaseOpModel() argument
36 axis_type_(axis_type), in ArgBaseOpModel()
40 if (axis_type == TensorType_INT64) { in ArgBaseOpModel()
42 AddConstInput(axis_type, {static_cast<int64_t>(axis_value)}, {1}); in ArgBaseOpModel()
44 axis_ = AddConstInput(axis_type, {axis_value}, {1}); in ArgBaseOpModel()
47 axis_ = AddInput(axis_type); in ArgBaseOpModel()
85 int axis_value, TensorType axis_type, bool constant_axis, in ArgMaxOpModel() argument
87 : ArgBaseOpModel(input_type, axis_value, axis_type, constant_axis, in ArgMaxOpModel()
100 int axis_value, TensorType axis_type, bool constant_axis, in ArgMinOpModel() argument
102 : ArgBaseOpModel(input_type, axis_value, axis_type, constant_axis, in ArgMinOpModel()
|
| D | arg_min_max.cc | 143 #define TF_LITE_ARG_MIN_MAX(data_type, axis_type, output_type) \ in Eval() argument
146 GetTensorData<axis_type>(axis), GetTensorShape(output), \ in Eval()
|
| /external/tensorflow/tensorflow/lite/micro/kernels/ |
| D | arg_min_max.cc | 56 #define TF_LITE_ARG_MIN_MAX(data_type, axis_type, output_type) \ in Eval() argument
59 tflite::micro::GetTensorData<axis_type>(axis), \ in Eval()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | gather_op.cc | 168 DataType axis_type = context->input_type(2); in XlaGatherWithBatchDimsOpImpl() local
169 if (axis_type != DT_INT32 && axis_type != DT_INT64) { in XlaGatherWithBatchDimsOpImpl()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/ |
| D | tf_ops_a_m.cc | 1176 auto axis_type = mlir::RankedTensorType::get({}, rewriter.getIntegerType(64)); in matchAndRewrite() local
1178 loc, DenseIntElementsAttr::get(axis_type, hoist_params->lhs_axis)); in matchAndRewrite()
1180 loc, DenseIntElementsAttr::get(axis_type, hoist_params->rhs_axis)); in matchAndRewrite()
|