Searched refs:layout_func (Results 1 – 2 of 2) sorted by relevance
73 const std::function<xla::Layout(const xla::Shape&)>& layout_func, in AssignLayout() argument77 if (!has_layout && layout_func) { in AssignLayout()78 layout = layout_func(*shape); in AssignLayout()132 const std::function<xla::Layout(const xla::Shape&)>& layout_func, in GetShapeWithLayout() argument156 layout_func, &shapes.back())); in GetShapeWithLayout()176 TF_RETURN_IF_ERROR(AssignLayout(minor_to_major, layout_func, output_shape)); in GetShapeWithLayout()
63 const std::function<xla::Layout(const xla::Shape&)>& layout_func,