Home
last modified time | relevance | path

Searched refs:index_shape (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dslicing.cc140 TF_ASSIGN_OR_RETURN(Shape index_shape, builder->GetShape(index)); in TorchGather()
141 ShapeUtil::AppendMajorDimension(1, &index_shape); in TorchGather()
147 to_concat.push_back(Reshape(index, index_shape.dimensions())); in TorchGather()
149 to_concat.push_back(Iota(builder, index_shape, i)); in TorchGather()
168 TF_ASSIGN_OR_RETURN(Shape index_shape, builder->GetShape(index)); in TorchIndexSelect()
177 gather_dnums.set_index_vector_dim(index_shape.rank()); in TorchIndexSelect()
/external/tensorflow/tensorflow/lite/kernels/
Dembedding_lookup_test.cc37 BaseEmbeddingLookupOpModel(std::initializer_list<int> index_shape, in BaseEmbeddingLookupOpModel() argument
44 BuildInterpreter({index_shape, weight_shape}); in BaseEmbeddingLookupOpModel()
80 HybridEmbeddingLookupOpModel(std::initializer_list<int> index_shape, in HybridEmbeddingLookupOpModel() argument
83 : BaseEmbeddingLookupOpModel(index_shape, weight_shape, type) {} in HybridEmbeddingLookupOpModel()
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Ddynamic_slice_ops.cc42 const TensorShape index_shape = ctx->InputShape("indices"); in Compile() local
46 TensorShapeUtils::IsVector(index_shape) && in Compile()
47 index_shape.num_elements() == input_shape.dims(), in Compile()
Dsplit_op.cc106 const TensorShape index_shape = ctx->InputShape(2); in Compile() local
/external/tensorflow/tensorflow/compiler/xla/service/
Dgather_expander.cc114 const Shape& index_shape = index_vector->shape(); in ExpandIndexVectorIntoOperandSpace() local
120 LiteralUtil::CreateFromDimensions(index_shape.element_type(), {0}))); in ExpandIndexVectorIntoOperandSpace()
125 LiteralUtil::CreateFromDimensions(index_shape.element_type(), {1}))); in ExpandIndexVectorIntoOperandSpace()
Dscatter_expander.cc136 const Shape& index_shape = index_vector->shape(); in ExpandIndexVectorIntoOperandSpace() local
141 LiteralUtil::CreateFromDimensions(index_shape.element_type(), {0}))); in ExpandIndexVectorIntoOperandSpace()
146 LiteralUtil::CreateFromDimensions(index_shape.element_type(), {1}))); in ExpandIndexVectorIntoOperandSpace()
Dshape_inference.cc2331 for (const Shape& index_shape : start_index_shapes) { in InferDynamicSliceShape() local
2332 if (!ShapeUtil::Compatible(first_index_shape, index_shape)) { in InferDynamicSliceShape()
2337 ShapeUtil::HumanString(index_shape)); in InferDynamicSliceShape()
2438 for (const Shape& index_shape : start_index_shapes) { in InferDynamicUpdateSliceShape() local
2439 if (!ShapeUtil::Compatible(first_index_shape, index_shape)) { in InferDynamicUpdateSliceShape()
2444 ShapeUtil::HumanString(index_shape)); in InferDynamicUpdateSliceShape()
/external/tensorflow/tensorflow/core/kernels/
Dbatch_kernels.cc595 TensorShape index_shape({0, 3}); in ProcessBatch() local
599 task.context->allocate_output(num_input_edges, index_shape, &output), in ProcessBatch()
633 const TensorShape index_shape({batch.num_tasks(), 3}); in EmitIndexTensor() local
636 context->allocate_output(output_index, index_shape, &index)); in EmitIndexTensor()
/external/tensorflow/tensorflow/lite/delegates/nnapi/
Dnnapi_delegate_test.cc3461 BaseEmbeddingLookupOpModel(std::initializer_list<int> index_shape, in BaseEmbeddingLookupOpModel() argument
3468 BuildInterpreter({index_shape, weight_shape}); in BaseEmbeddingLookupOpModel()