| /external/tensorflow/tensorflow/compiler/xla/service/ |
| hlo_element_type_converter.cc | 70 Shape subshape = ShapeUtil::GetTupleElementShape(shape, i); local
71 CHECK(!ShapeUtil::IsTuple(subshape));
72 if (subshape.element_type() == from_type) {
73 subshape = ShapeUtil::ChangeElementType(subshape, to_type);
75 new_tuple_subshapes.push_back(subshape);
|
| bfloat16_normalization.cc | 196 auto subshape = ShapeUtil::GetMutableSubshape(crs->mutable_shape(), {i}); local
199 HloInstruction::CreateGetTupleElement(*subshape, crs, i));
202 subshape->set_element_type(F32);
204 HloInstruction::CreateGetTupleElement(*subshape, crs, i));
207 ShapeUtil::ChangeElementType(*subshape, BF16), gte));
|
| /external/tensorflow/tensorflow/core/ops/ |
| math_ops.cc | 89 TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_batch_dims));
90 TF_RETURN_IF_ERROR(c->Subshape(b_shape, 0, -2, &b_batch_dims));
865 ShapeHandle subshape; local
866 TF_RETURN_IF_ERROR(c->Subshape(data_shape, 1, &subshape));
870 c->Concatenate(c->Vector(InferenceContext::kUnknownDim), subshape, &out));
889 ShapeHandle subshape; local
890 TF_RETURN_IF_ERROR(c->Subshape(data_shape, 1, &subshape));
894 c->Concatenate(c->Vector(InferenceContext::kUnknownDim), subshape, &out))
913 ShapeHandle subshape; local
954 ShapeHandle subshape; local
[all...] |
| /external/tensorflow/tensorflow/compiler/xla/ |
| literal_util.h | 97 // given NativeT (e.g., float). CHECKs if the subshape of the literal at the
114 // at the given shape index. CHECKs if the subshape of the literal at the
234 // literal rooted at 'dest_shape_index'. The subshape of this literal rooted
235 // at 'dest_shape_index' must be compatible with the subshape of 'src_literal'
241 // Similar to CopyFrom, but with move semantincs. The subshape of this literal
244 // allocated in this literal for the subshape at dest_shape_index is
657 // A data structure representing a subshape at a particular ShapeIndex within
694 // Gets or sets the subshape of this piece. This reference points to a
695 // subshape within the shape in the containing Literal (Literal::shape_).
696 const Shape& subshape() const { return *subshape_; function in class:xla::Literal::Piece
1127 const Shape& subshape = p.subshape(); local
[all...] |
| literal_util.cc | 97 const Shape& subshape = piece.subshape(); local
98 if (ShapeUtil::IsArray(subshape)) {
101 if (LayoutUtil::IsSparseArray(subshape)) {
103 LayoutUtil::MaxSparseElements(subshape.layout()),
104 ShapeUtil::Rank(subshape)));
130 // We need to iterate through the pieces to set the subshape pointer
148 // We need to iterate through the pieces to set the subshape pointer
167 if (ShapeUtil::IsArray(piece.subshape())) {
317 if (ShapeUtil::Equal(subshape(), src.subshape()))
709 Shape* subshape = ShapeUtil::GetMutableSubshape(&new_shape, shape_index); local
884 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); local
925 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); local
1076 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); local
[all...] |
