Searched defs:infeed_shape (Results 1 – 7 of 7) sorted by relevance
121 const auto infeed_shape = ShapeUtil::MakeShape(F32, {3}); in TEST_F() local223 const auto build_body = [&result_shape](const Shape& infeed_shape) { in TEST_F()
42 xla::Shape infeed_shape = ApiConverter::FromC(&c_infeed_shape); in GetTPUInfeedLayout() local
265 if (absl::optional<Shape> infeed_shape = GetXfeedShape( in ReplayComputation() local
65 xla::Shape infeed_shape = ApiConverter::FromC(&c_infeed_shape); in GetTPUInfeedLayout() local
1167 const Shape& infeed_shape() const { in infeed_shape() function
1109 const Shape& infeed_shape, HloInstruction* token_operand, in CreateInfeed()
2047 HloInfeedInstruction::HloInfeedInstruction(const Shape& infeed_shape, in HloInfeedInstruction()