Searched refs:indices_type (Results 1 – 9 of 9) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | unsorted_segment_join_op.cc | 152 #define REGISTER_CPU_KERNEL(indices_type, num_segments_type) \ argument
156 .TypeConstraint<indices_type>("Tindices") \
158 UnsortedSegmentJoinOp<indices_type, num_segments_type>);
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | softmax_op.cc | 197 const DataType indices_type = input_type(1); in Compile() local
216 xla::And(xla::Le(XlaHelpers::Zero(builder, indices_type), indices), in Compile()
218 builder, indices_type, depth))), in Compile()
|
| D | gather_op.cc | 273 DataType indices_type = context->input_type(1); in Compile() local
297 indices_type, builder, &gather)); in Compile()
|
| D | tensor_list_ops.cc | 332 DataType indices_type = ctx->input_type(1); in Compile() local
351 /*indices_are_nd=*/false, dtype_, indices_type, in Compile()
|
| /external/tensorflow/tensorflow/lite/kernels/ |
| D | one_hot_test.cc | 34 T off_value = 0, TensorType indices_type = TensorType_INT32) { in OneHotOpModel() argument
35 indices_ = AddInput(indices_type); in OneHotOpModel()
|
| /external/tensorflow/tensorflow/compiler/mlir/lite/ir/ |
| D | tfl_ops.cc | 614 auto indices_type = indices.getType().cast<TensorType>(); in BuildGatherOp() local
617 if (!params_type.hasRank() || !indices_type.hasRank()) in BuildGatherOp()
623 int64_t indices_rank = indices_type.getRank(); in BuildGatherOp()
648 std::copy(std::begin(indices_type.getShape()), in BuildGatherOp()
649 std::end(indices_type.getShape()), std::begin(shape) + axis_i); in BuildGatherOp()
659 std::copy(std::begin(indices_type.getShape()), in BuildGatherOp()
660 std::end(indices_type.getShape()), std::begin(shape) + axis_i); in BuildGatherOp()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | algebraic_simplifier.cc | 1605 auto indices_type = dynamic_slice->operand(1)->shape().element_type(); in OptimizeDotOfGather() local
1606 Shape s_shape = ShapeUtil::MakeShape(indices_type, {1}); in OptimizeDotOfGather()
1608 Shape d_shape = ShapeUtil::MakeShape(indices_type, {2}); in OptimizeDotOfGather()
|
| /external/tensorflow/tensorflow/lite/toco/ |
| D | export_tensorflow.cc | 1798 const tensorflow::DataType indices_type = in ConvertReduceOperator() local
1800 (*new_op->mutable_attr())["Tidx"].set_type(indices_type); in ConvertReduceOperator()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 3293 auto indices_type = in matchAndRewrite() local
3296 op.getLoc(), indices_type, rewriter.getI64IntegerAttr(first_dim_size)); in matchAndRewrite()
|