Home
last modified time | relevance | path

Searched refs:scatter_indices_shape (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dscatter_expander.cc33 const Shape& scatter_indices_shape = scatter_indices->shape(); in TransposeIndexVectorDimToLast() local
35 if (scatter_indices_shape.dimensions_size() == index_vector_dim) { in TransposeIndexVectorDimToLast()
39 if (index_vector_dim == (scatter_indices_shape.dimensions_size() - 1)) { in TransposeIndexVectorDimToLast()
44 permutation.reserve(scatter_indices_shape.dimensions_size()); in TransposeIndexVectorDimToLast()
45 for (int64 i = 0, e = scatter_indices_shape.dimensions_size(); i < e; i++) { in TransposeIndexVectorDimToLast()
115 const Shape& scatter_indices_shape, HloInstruction* updates, in AdjustScatterDims() argument
117 int64 num_scatter_dims = scatter_indices_shape.dimensions_size(); in AdjustScatterDims()
118 if (index_vector_dim < scatter_indices_shape.dimensions_size()) { in AdjustScatterDims()
361 const Shape& scatter_indices_shape = scatter_indices->shape(); in ExpandScatter() local
363 for (int64 i = 0, e = scatter_indices_shape.dimensions_size(); i < e; i++) { in ExpandScatter()
[all …]
Dshape_inference.cc3070 const Shape& operand_shape, absl::Span<const int64> scatter_indices_shape, in ValidateScatterDimensionNumbers() argument
3126 scatter_indices_shape[dim_numbers.index_vector_dim()]) { in ValidateScatterDimensionNumbers()
3133 scatter_indices_shape[dim_numbers.index_vector_dim()]); in ValidateScatterDimensionNumbers()
3164 const Shape& operand_shape, const Shape& scatter_indices_shape, in InferScatterShape() argument
3170 ExpectArray(scatter_indices_shape, "scatter indices of scatter op")); in InferScatterShape()
3173 if (!ShapeUtil::ElementIsIntegral(scatter_indices_shape)) { in InferScatterShape()
3176 ShapeUtil::HumanString(scatter_indices_shape)); in InferScatterShape()
3179 if (scatter_indices_shape.dimensions_size() < in InferScatterShape()
3186 scatter_indices_shape.dimensions_size(), in InferScatterShape()
3198 ArraySliceToVector(AsInt64Slice(scatter_indices_shape.dimensions())); in InferScatterShape()
Dshape_inference.h291 const Shape& operand_shape, const Shape& scatter_indices_shape,
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.cc1123 Shape scatter_indices_shape = scatter_indices->shape(); in EmitScatter() local
1124 if (dim_numbers.index_vector_dim() == scatter_indices_shape.rank()) { in EmitScatter()
1125 scatter_indices_shape.add_dimensions(1); in EmitScatter()
1126 scatter_indices_shape.mutable_layout()->add_minor_to_major( in EmitScatter()
1145 raw_scatter_index_multidim, scatter_indices_shape, index.GetType()); in EmitScatter()
1151 scatter_indices_shape, scatter_indices->shape(), &b_))); in EmitScatter()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.cc1872 TF_ASSIGN_OR_RETURN(const Shape& scatter_indices_shape, in Scatter()
1879 input_shape, scatter_indices_shape, updates_shape, in Scatter()