Searched refs:shape_mask (Results 1 – 2 of 2) sorted by relevance
2236 HloInstruction* shape_mask = MakeBroadcastHlo( in Propagate() local2242 MakeSelectHlo(shape_mask, new_consumer, final_selection)); in Propagate()2355 HloInstruction* shape_mask = MakeBroadcastHlo( in SelectValidPortion() local2358 VLOG(1) << "Shape mask made " << shape_mask->ToString(); in SelectValidPortion()2363 TF_ASSIGN_OR_RETURN(new_instr, MakeSelectHlo(shape_mask, new_instr, zeroes)); in SelectValidPortion()
1882 shape_mask = mask.get_shape()1883 ndims_mask = shape_mask.ndims1895 shape_tensor[axis:axis + ndims_mask].assert_is_compatible_with(shape_mask)