Searched refs:index_shape (Results 1 – 9 of 9) sorted by relevance
/external/tensorflow/tensorflow/compiler/xla/client/lib/ |
D | slicing.cc | 140 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/ |
D | embedding_lookup_test.cc | 37 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/ |
D | dynamic_slice_ops.cc | 42 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()
|
D | split_op.cc | 106 const TensorShape index_shape = ctx->InputShape(2); in Compile() local
|
/external/tensorflow/tensorflow/compiler/xla/service/ |
D | gather_expander.cc | 114 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()
|
D | scatter_expander.cc | 136 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()
|
D | shape_inference.cc | 2331 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/ |
D | batch_kernels.cc | 595 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/ |
D | nnapi_delegate_test.cc | 3461 BaseEmbeddingLookupOpModel(std::initializer_list<int> index_shape, in BaseEmbeddingLookupOpModel() argument 3468 BuildInterpreter({index_shape, weight_shape}); in BaseEmbeddingLookupOpModel()
|