Home
last modified time | relevance | path

Searched refs:start_indices_shape (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dgather_expander.cc31 const Shape& start_indices_shape = start_indices->shape(); in TransposeIndexVectorDimToLast() local
33 if (start_indices_shape.dimensions_size() == index_vector_dim) { in TransposeIndexVectorDimToLast()
37 if (index_vector_dim == (start_indices_shape.dimensions_size() - 1)) { in TransposeIndexVectorDimToLast()
42 permutation.reserve(start_indices_shape.dimensions_size()); in TransposeIndexVectorDimToLast()
43 for (int64 i = 0, e = start_indices_shape.dimensions_size(); i < e; i++) { in TransposeIndexVectorDimToLast()
87 const Shape& start_indices_shape, HloInstruction* accumulator, in AdjustBatchDimsInAccumulator() argument
90 batch_dim_bounds.reserve(start_indices_shape.dimensions_size()); in AdjustBatchDimsInAccumulator()
91 for (int64 i = 0, e = start_indices_shape.dimensions_size(); i < e; i++) { in AdjustBatchDimsInAccumulator()
93 batch_dim_bounds.push_back(start_indices_shape.dimensions(i)); in AdjustBatchDimsInAccumulator()
275 const Shape& start_indices_shape = start_indices->shape(); in GatherLoopTripCount() local
[all …]
Dshape_inference.cc2487 const Shape& start_indices_shape = start_index_shapes[0]; in InferDynamicSliceShape() local
2491 ShapeUtil::HumanString(start_indices_shape), in InferDynamicSliceShape()
2495 ExpectArray(start_indices_shape, "start indices of dynamic slice")); in InferDynamicSliceShape()
2497 if (start_indices_shape.rank() != 1) { in InferDynamicSliceShape()
2500 start_indices_shape.rank()); in InferDynamicSliceShape()
2503 if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { in InferDynamicSliceShape()
2508 const int64 start_num_dims = start_indices_shape.dimensions(0); in InferDynamicSliceShape()
2513 start_num_dims, ShapeUtil::HumanString(start_indices_shape), in InferDynamicSliceShape()
2592 const Shape& start_indices_shape = start_index_shapes[0]; in InferDynamicUpdateSliceShape() local
2593 TF_RETURN_IF_ERROR(ExpectArray(start_indices_shape, in InferDynamicUpdateSliceShape()
[all …]
Dshape_inference.h314 const Shape& input_shape, const Shape& start_indices_shape,
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Ddynamic_slice_ops.cc87 const TensorShape start_indices_shape = ctx->InputShape("start_indices"); in Compile() local
92 TensorShapeUtils::IsVector(start_indices_shape) && in Compile()
93 start_indices_shape.num_elements() == rank, in Compile()
98 start_indices_shape.DebugString())); in Compile()
/external/tensorflow/tensorflow/compiler/xla/tests/
Ddynamic_ops_test.cc775 auto start_indices_shape = ShapeUtil::MakeShape(S32, {}); in BM_DynamicSlice() local
781 Parameter(&builder, i, start_indices_shape, "start_indices"); in BM_DynamicSlice()
787 ->AllocateScopedShapedBuffer(start_indices_shape, &allocator, in BM_DynamicSlice()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.cc2276 TF_ASSIGN_OR_RETURN(const Shape* start_indices_shape, in Gather()
2279 *input_shape, *start_indices_shape, in Gather()