Home
last modified time | relevance | path

Searched refs:tuple_shapes (Results 1 – 25 of 46) sorted by relevance

12

/external/tensorflow/tensorflow/python/tpu/
Dtpu_infeed_test.py36 self.assertEqual(i.tuple_shapes, None)
43 self.assertEqual(i.tuple_shapes, None)
45 i = tpu_feed.InfeedQueue(tuple_shapes=[[1], [2, 3]])
48 self.assertEqual(i.tuple_shapes, [[1], [2, 3]])
53 self.assertEqual(i.tuple_shapes, None)
62 i = tpu_feed.InfeedQueue(number_of_tuple_elements=2, tuple_shapes=[[1]])
66 i = tpu_feed.InfeedQueue(tuple_shapes=[[1], [2, 3]], shard_dimensions=[1])
78 self.assertEqual(i.tuple_shapes, [[1], [2, 3]])
80 self.assertEqual(i.tuple_shapes, [[1, 2], [3, 4]])
90 self.assertEqual(i.tuple_shapes, [[6], [3, 18]])
[all …]
Dtpu_feed.py131 tuple_shapes=None, argument
167 elif tuple_shapes is not None:
168 number_of_tuple_elements = len(tuple_shapes)
187 if tuple_shapes is not None:
188 self.set_tuple_shapes(tuple_shapes)
201 if self.tuple_shapes is not None:
251 def tuple_shapes(self): member in InfeedQueue
255 def set_tuple_shapes(self, tuple_shapes): argument
271 if len(tuple_shapes) != self.number_of_tuple_elements:
273 (str(tuple_shapes), self.number_of_tuple_elements))
[all …]
Dtpu_test.py78 tuple_shapes=[inputs.shape, scale.shape])
/external/tensorflow/tensorflow/compiler/xla/
Dlayout_util_test.cc270 EXPECT_TRUE(shape.tuple_shapes(0).has_layout()); in TEST_F()
271 EXPECT_TRUE(shape.tuple_shapes(2).tuple_shapes(1).has_layout()); in TEST_F()
276 EXPECT_FALSE(shape.tuple_shapes(0).has_layout()); in TEST_F()
277 EXPECT_FALSE(shape.tuple_shapes(2).tuple_shapes(1).has_layout()); in TEST_F()
297 EXPECT_FALSE(LayoutUtil::Equal(shape.tuple_shapes(0).layout(), in TEST_F()
298 shape.tuple_shapes(1).layout())); in TEST_F()
300 EXPECT_TRUE(LayoutUtil::Equal(shape.tuple_shapes(0).layout(), in TEST_F()
301 shape.tuple_shapes(1).layout())); in TEST_F()
303 LayoutUtil::GetDefaultLayoutForShape(shape.tuple_shapes(0)), in TEST_F()
304 shape.tuple_shapes(1).layout())); in TEST_F()
Dshape.cc50 for (const ShapeProto& element_shape : shape_proto.tuple_shapes()) { in Shape()
69 for (const Shape& shape : tuple_shapes()) { in ToProto()
123 lhs.tuple_shapes(), rhs.tuple_shapes(), in operator ()()
Dshape_util.cc330 absl::c_any_of(shape.tuple_shapes(), in IsNestedTuple()
347 TF_DCHECK_OK(ValidateShapeWithOptionalLayout(shape.tuple_shapes(index))); in GetTupleElementShape()
348 return shape.tuple_shapes(index); in GetTupleElementShape()
365 std::vector<Shape> new_elements(tuple.tuple_shapes().begin() + start, in SliceTuple()
366 tuple.tuple_shapes().begin() + limit); in SliceTuple()
395 for (const Shape& element_shape : shape.tuple_shapes()) { in ElementsInRecursive()
406 for (const Shape& element_shape : shape.tuple_shapes()) { in HasPrimitiveType()
427 for (const Shape& elem_shape : shape.tuple_shapes()) { in HumanString()
451 for (const Shape& elem_shape : shape.tuple_shapes()) { in HumanStringWithLayout()
632 for (auto& element_shape : shape.tuple_shapes()) { in ValidateShapeWithOptionalLayoutInternal()
[all …]
Dlayout_util.cc180 for (auto& element_shape : shape.tuple_shapes()) { in ValidateLayoutInShape()
309 return absl::c_all_of(shape.tuple_shapes(), in HasLayout()
388 TF_RETURN_IF_ERROR(CopyLayoutInternal(src.tuple_shapes(i), in CopyLayoutInternal()
421 if (!LayoutsInShapesEqual(lhs.tuple_shapes(i), rhs.tuple_shapes(i))) { in LayoutsInShapesEqual()
Dshape.h113 const Shape& tuple_shapes(int index) const { return tuple_shapes_.at(index); } in tuple_shapes() function
120 const std::vector<Shape>& tuple_shapes() const { return tuple_shapes_; } in tuple_shapes() function
Dshape_tree.h401 current_count += CountSubshapes(shape.tuple_shapes(i)); in CountSubshapes()
434 InitChildren(shape.tuple_shapes(i), init_value, &nodes_.back(), in InitChildren()
470 InitChildren(shape.tuple_shapes(i), &nodes_.back(), in InitChildren()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dinfeed_thunk.cc43 TF_RET_CHECK(infeed_shape.tuple_shapes().size() == 2) in ExecuteOnStream()
45 TF_RET_CHECK(infeed_shape.tuple_shapes(1).IsToken()) in ExecuteOnStream()
48 ShapeUtil::Equal(infeed_buffers.shape(), infeed_shape.tuple_shapes(0))) in ExecuteOnStream()
50 << ShapeUtil::HumanStringWithLayout(infeed_shape.tuple_shapes(0)) in ExecuteOnStream()
Dcudnn_batchnorm_thunk.cc156 CHECK(LayoutUtil::LayoutsInShapesEqual(hlo->shape().tuple_shapes(0), in CudnnBatchNormForwardTrainingThunk()
158 for (const auto& tuple_shape : hlo->shape().tuple_shapes()) { in CudnnBatchNormForwardTrainingThunk()
172 hlo_instruction()->shape().tuple_shapes(0), feature_index_); in ExecuteOnStream()
238 CHECK(LayoutUtil::LayoutsInShapesEqual(hlo->shape().tuple_shapes(0), in CudnnBatchNormBackwardThunk()
240 CHECK(LayoutUtil::LayoutsInShapesEqual(hlo->shape().tuple_shapes(0), in CudnnBatchNormBackwardThunk()
242 for (const auto& tuple_shape : hlo->shape().tuple_shapes()) { in CudnnBatchNormBackwardThunk()
257 hlo_instruction()->shape().tuple_shapes(0), feature_index_); in ExecuteOnStream()
Dcudnn_conv_pad_for_tensor_cores.cc77 CHECK_EQ(0, conv->shape().tuple_shapes(1).dimensions(0)) in PadConv()
85 const Shape& result_shape = conv->shape().tuple_shapes(0); in PadConv()
126 const Shape& result_shape = conv->shape().tuple_shapes(0); in PadForTensorCores()
141 Shape new_result_shape = conv->shape().tuple_shapes(0); in PadForTensorCores()
Dcudnn_batchnorm_rewriter.cc128 libcall->shape().tuple_shapes(2), libcall, 2)); in HandleBatchNormTraining()
147 libcall->shape().tuple_shapes(0), libcall, 0)), in HandleBatchNormTraining()
149 libcall->shape().tuple_shapes(1), libcall, 1)), in HandleBatchNormTraining()
Dcudnn_conv_algorithm_picker.cc143 instr->shape().tuple_shapes(0).element_type() == xla::F16; in PickBestAlgorithm()
222 &stream, ShapeUtil::ByteSizeOf(instr->shape().tuple_shapes(0)))); in PickBestAlgorithm()
380 {instr->shape().tuple_shapes(0), in RunOnInstruction()
401 new_call_shape.tuple_shapes(0), new_call, 0)), in RunOnInstruction()
Dgpu_layout_assignment_test.cc246 batchnorm->shape().tuple_shapes(0).layout())) in TEST_F()
328 batchnorm->shape().tuple_shapes(0).layout())) in TEST_F()
332 batchnorm->shape().tuple_shapes(0).layout())) in TEST_F()
Dcudnn_fused_conv_rewriter.cc210 int64 num_output_feature = conv->shape().tuple_shapes(0).dimensions( in TryRewriteToCudnnForwardRelu()
240 return HloInstruction::CreateGetTupleElement(conv->shape().tuple_shapes(0), in TryRewriteToCudnnForwardRelu()
Dcudnn_conv_padding_legalization.cc287 Shape backward_conv_shape = backward_conv->shape().tuple_shapes(0); in CanonicalizeBackwardInputConvolution()
341 new_backward_conv_call->shape().tuple_shapes(1), in CanonicalizeBackwardInputConvolution()
Dcudnn_conv_runner.cc331 const auto& conv_result_shape = conv->shape().tuple_shapes(0); in GetCudnnConvParams()
417 conv->shape().tuple_shapes(0).element_type(); in RunCudnnConv()
/external/tensorflow/tensorflow/compiler/xla/service/
Dwhile_loop_simplifier.cc174 while_init->shape().tuple_shapes(old_idx)); in TryRemoveDeadWhileParams()
257 while_init->shape().tuple_shapes(old_idx), while_init, old_idx))); in TryRemoveDeadWhileParams()
293 new_while_op->shape().tuple_shapes(gte_idx), new_while_op, in TryRemoveDeadWhileParams()
298 while_init->shape().tuple_shapes(old_idx), while_init, old_idx))); in TryRemoveDeadWhileParams()
348 new_while_shape_elems.push_back(while_shape.tuple_shapes(i)); in TryRemoveConstantParams()
371 while_shape.tuple_shapes(i), instr, i))); in TryRemoveConstantParams()
388 while_shape.tuple_shapes(i), instr, j))); in TryRemoveConstantParams()
593 const Shape& subshape = desired_shape.tuple_shapes(i); in UnflattenTupleInstr()
611 desired_shape.tuple_shapes(i), new_instrs); in UnflattenTupleInstr()
631 const Shape& subshape = shape.tuple_shapes(i); in GetFlatTupleElems()
[all …]
Dtuple_util.cc34 input_shape.tuple_shapes(i), input_tuple, i))); in ExtractPrefix()
53 input_shape.tuple_shapes(i), input_tuple, i))); in AppendSuffix()
Dhlo_module_dce.cc74 while_body_param->shape().tuple_shapes(i), while_body_param, in RunWhileDCE()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dtesting.cc42 for (const Shape& s : shape.tuple_shapes()) { in DataSizeOfShape()
56 for (const Shape& s : shape.tuple_shapes()) { in BuildFakeDataOpOnDevice()
/external/tensorflow/tensorflow/compiler/xla/python_api/
Dxla_shape.py62 component_message = self.message.tuple_shapes.add()
82 def tuple_shapes(self): member in Shape
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client.py478 def tuple_shape(tuple_shapes): argument
480 if (not isinstance(tuple_shapes, (tuple, list)) or
481 not all(isinstance(t, Shape) for t in tuple_shapes)):
483 return Shape(tuple_shapes, tuple)
536 def tuple_shapes(self): member in Shape
578 children = tuple(child.map_leaves(f) for child in self.tuple_shapes())
616 for shape in self.tuple_shapes():
617 shape.serialize(proto.tuple_shapes.add())
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dwhile_op.cc162 body.xla_output_shape.tuple_shapes(i); in VerifyBodyInputAndOutputShapeMatch()
168 body.xla_output_shape.tuple_shapes(ctx->num_inputs()); in VerifyBodyInputAndOutputShapeMatch()
395 arg.shape = body_output_shape.tuple_shapes(i); in Compile()
491 GetTensorListBufferShape(body_input_shape.tuple_shapes(i), in Compile()

12