Searched defs:loss_layout (Results 1 – 1 of 1) sorted by relevance
373 const Layout& labels_layout, const Layout& loss_layout, in MaybeRelayoutInputs()469 const Layout& output_layout, const Layout& loss_layout, in MaybeRelayoutOutputs()669 const Layout loss_layout = backprop_layout.Truncate(/*split_point=*/1); in ComputeLayoutForward() local681 absl::optional<Layout> loss_layout; in ComputeLayoutBackward() local