Searched refs:flatten_shape (Results 1 – 1 of 1) sorted by relevance
599 def flatten_shape(s: XlaShape) -> Sequence[Tuple[Sequence[int], XlaShape]]: function740 for param_index, element in flatten_shape(c.GetShape(arg)):752 for output_index, element in flatten_shape(c.GetShape(out_tuple)):