Searched refs:desired_shape (Results 1 – 2 of 2) sorted by relevance
795 absl::Span<HloInstruction*> instrs, const Shape& desired_shape, in UnflattenTupleInstr() argument797 CHECK(desired_shape.IsTuple()) << ShapeUtil::HumanString(desired_shape); in UnflattenTupleInstr()804 for (int64 i = 0; i < desired_shape.tuple_shapes_size(); ++i) { in UnflattenTupleInstr()805 const Shape& subshape = desired_shape.tuple_shapes(i); in UnflattenTupleInstr()823 desired_shape.tuple_shapes(i), new_instrs); in UnflattenTupleInstr()
1588 PartialTensorShape desired_shape; in __anon42d741192302() local1589 TF_RETURN_IF_ERROR(c->GetAttr("shape", &desired_shape)); in __anon42d741192302()1591 int rank = desired_shape.dims(); in __anon42d741192302()1596 desired_shape, &desired_shape_handle)); in __anon42d741192302()