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()
314 const Shape& start_indices_shape = start_indices->shape(); in ExpandInstruction() local
[all …]
Dshape_inference.cc2279 const Shape& start_indices_shape = start_index_shapes[0]; in InferDynamicSliceShape() local
2283 ShapeUtil::HumanString(start_indices_shape), in InferDynamicSliceShape()
2287 ExpectArray(start_indices_shape, "start indices of dynamic slice")); in InferDynamicSliceShape()
2289 if (start_indices_shape.rank() != 1) { in InferDynamicSliceShape()
2292 start_indices_shape.rank()); in InferDynamicSliceShape()
2295 if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { in InferDynamicSliceShape()
2300 const int64 start_num_dims = start_indices_shape.dimensions(0); in InferDynamicSliceShape()
2305 start_num_dims, ShapeUtil::HumanString(start_indices_shape), in InferDynamicSliceShape()
2384 const Shape& start_indices_shape = start_index_shapes[0]; in InferDynamicUpdateSliceShape() local
2385 TF_RETURN_IF_ERROR(ExpectArray(start_indices_shape, in InferDynamicUpdateSliceShape()
[all …]
Dshape_inference.h283 const Shape& input_shape, const Shape& start_indices_shape,
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Ddynamic_slice_ops.cc76 const TensorShape start_indices_shape = ctx->InputShape("start_indices"); in Compile() local
80 TensorShapeUtils::IsVector(start_indices_shape) && in Compile()
81 start_indices_shape.num_elements() == input_shape.dims(), in Compile()
86 start_indices_shape.DebugString())); in Compile()
/external/tensorflow/tensorflow/compiler/xla/tests/
Ddynamic_ops_test.cc758 auto start_indices_shape = ShapeUtil::MakeShape(S32, {}); in BM_DynamicSlice() local
764 Parameter(&builder, i, start_indices_shape, "start_indices"); in BM_DynamicSlice()
770 ->AllocateScopedShapedBuffer(start_indices_shape, &allocator, in BM_DynamicSlice()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.cc764 TF_ASSIGN_OR_RETURN(const Shape& start_indices_shape, in DynamicSlice()
768 operand_shape, {start_indices_shape}, slice_sizes)); in DynamicSlice()
815 TF_ASSIGN_OR_RETURN(const Shape& start_indices_shape, in DynamicUpdateSlice()
819 operand_shape, update_shape, {start_indices_shape})); in DynamicUpdateSlice()
1847 TF_ASSIGN_OR_RETURN(const Shape& start_indices_shape, in Gather()
1850 input_shape, start_indices_shape, in Gather()