Searched refs:start_indices_shape (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | gather_expander.cc | 33 const Shape& start_indices_shape = start_indices->shape(); in TransposeIndexVectorDimToLast() local
35 if (start_indices_shape.dimensions_size() == index_vector_dim) { in TransposeIndexVectorDimToLast()
39 if (index_vector_dim == (start_indices_shape.dimensions_size() - 1)) { in TransposeIndexVectorDimToLast()
44 permutation.reserve(start_indices_shape.dimensions_size()); in TransposeIndexVectorDimToLast()
45 for (int64_t i = 0, e = start_indices_shape.dimensions_size(); i < e; i++) { in TransposeIndexVectorDimToLast()
89 const Shape& start_indices_shape, HloInstruction* accumulator, in AdjustBatchDimsInAccumulator() argument
92 batch_dim_bounds.reserve(start_indices_shape.dimensions_size()); in AdjustBatchDimsInAccumulator()
93 for (int64_t i = 0, e = start_indices_shape.dimensions_size(); i < e; i++) { in AdjustBatchDimsInAccumulator()
95 batch_dim_bounds.push_back(start_indices_shape.dimensions(i)); in AdjustBatchDimsInAccumulator()
277 const Shape& start_indices_shape = start_indices->shape(); in GatherLoopTripCount() local
[all …]
|
| D | shape_inference.cc | 2604 const Shape& start_indices_shape = start_index_shapes[0]; in InferDynamicSliceShape() local
2608 ShapeUtil::HumanString(start_indices_shape), in InferDynamicSliceShape()
2612 ExpectArray(start_indices_shape, "start indices of dynamic slice")); in InferDynamicSliceShape()
2614 if (start_indices_shape.rank() != 1) { in InferDynamicSliceShape()
2617 start_indices_shape.rank()); in InferDynamicSliceShape()
2620 if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { in InferDynamicSliceShape()
2625 const int64_t start_num_dims = start_indices_shape.dimensions(0); in InferDynamicSliceShape()
2630 start_num_dims, ShapeUtil::HumanString(start_indices_shape), in InferDynamicSliceShape()
2709 const Shape& start_indices_shape = start_index_shapes[0]; in InferDynamicUpdateSliceShape() local
2710 TF_RETURN_IF_ERROR(ExpectArray(start_indices_shape, in InferDynamicUpdateSliceShape()
[all …]
|
| D | shape_inference.h | 345 const Shape& input_shape, const Shape& start_indices_shape,
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | dynamic_ops_test.cc | 778 auto start_indices_shape = ShapeUtil::MakeShape(S32, {}); in BM_DynamicSlice() local
784 Parameter(&builder, i, start_indices_shape, "start_indices"); in BM_DynamicSlice()
790 ->AllocateScopedShapedBuffer(start_indices_shape, &allocator, in BM_DynamicSlice()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 2401 TF_ASSIGN_OR_RETURN(const Shape* start_indices_shape, in Gather()
2404 *input_shape, *start_indices_shape, in Gather()
|