Searched refs:pred_shape (Results 1 – 2 of 2) sorted by relevance
155 const Shape pred_shape = in Run() local165 pred_shape, iota, broadcasted_effective_size, in Run()
421 auto pred_shape = ShapeUtil::MakeShape(PRED, {8, 7}); in TEST_F() local433 builder.AddInstruction(HloInstruction::CreateReshape(pred_shape, pred)); in TEST_F()455 auto pred_shape = ShapeUtil::MakeShape(PRED, {}); in TEST_F() local463 builder.AddInstruction(HloInstruction::CreateReshape(pred_shape, pred)); in TEST_F()