Lines Matching refs:builder
59 ComputationBuilder builder(client_, "condition");
60 auto prev = builder.Parameter(0, result_shape, "prev");
61 builder.Gt(builder.ConstantR0<int32>(5), prev);
62 condition = builder.Build().ConsumeValueOrDie();
68 ComputationBuilder builder(client_, "body");
69 auto prev = builder.Parameter(0, result_shape, "prev");
70 auto input = builder.ConstantR0<int32>(1);
71 auto result = builder.Add(input, prev);
72 body = builder.Build().ConsumeValueOrDie();
76 ComputationBuilder builder(client_, TestName());
77 auto init = builder.ConstantR0<int32>(0);
78 auto result = builder.While(condition, body, init);
79 auto shape = builder.GetShape(result).ConsumeValueOrDie();
81 ComputeAndCompareR0<int32>(&builder, 5, {});
96 ComputationBuilder builder(client_, "condition");
97 auto prev = builder.Parameter(0, result_shape, "prev");
98 builder.Gt(builder.ConstantR0<int64>(5), prev);
99 condition = builder.Build().ConsumeValueOrDie();
105 ComputationBuilder builder(client_, "body");
106 auto prev = builder.Parameter(0, result_shape, "prev");
107 auto input = builder.ConstantR0<int64>(1);
108 auto result = builder.Add(input, prev);
109 body = builder.Build().ConsumeValueOrDie();
113 ComputationBuilder builder(client_, TestName());
114 auto init = builder.ConstantR0<int64>(0);
115 auto result = builder.While(condition, body, init);
116 auto shape = builder.GetShape(result).ConsumeValueOrDie();
118 ComputeAndCompareR0<int64>(&builder, 5, {});
128 ComputationBuilder builder(client_, "condition");
129 auto prev = builder.Parameter(0, result_shape, "prev");
130 builder.Gt(builder.ConstantR0<int32>(5), prev);
131 condition = builder.Build().ConsumeValueOrDie();
137 ComputationBuilder builder(client_, "body");
138 auto prev = builder.Parameter(0, result_shape, "prev");
139 auto input = builder.ConstantR0<int32>(1);
140 auto result = builder.Add(input, prev);
141 body = builder.Build().ConsumeValueOrDie();
145 ComputationBuilder builder(client_, TestName());
146 auto init = builder.Reduce(builder.ConstantR1<int32>(2, 1),
147 builder.ConstantR0<int32>(0),
148 CreateScalarAddComputation(S32, &builder), {0});
149 auto result = builder.While(condition, body, init);
150 auto shape = builder.GetShape(result).ConsumeValueOrDie();
152 ComputeAndCompareR0<int32>(&builder, 5, {});
161 ComputationBuilder builder(client_, "condition");
162 auto prev = builder.Parameter(0, result_shape, "prev");
163 builder.Ne(builder.ConstantR0<bool>(true), prev);
164 condition = builder.Build().ConsumeValueOrDie();
170 ComputationBuilder builder(client_, "body");
171 auto prev = builder.Parameter(0, result_shape, "prev");
172 auto result = builder.Or(prev, builder.ConstantR0<bool>(true));
173 body = builder.Build().ConsumeValueOrDie();
177 ComputationBuilder builder(client_, TestName());
178 auto init = builder.Ne(builder.ConstantR0<bool>(false),
179 builder.ConstantR0<bool>(true));
180 auto result = builder.While(condition, body, init);
182 ComputeAndCompareR0<bool>(&builder, true, {});
199 ComputationBuilder builder(client_, "add");
200 auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
201 auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y");
202 builder.Add(x, y);
203 add = builder.Build().ConsumeValueOrDie();
210 ComputationBuilder builder(client_, "condition");
211 auto prev = builder.Parameter(0, result_shape, "prev");
212 auto sum = builder.Reduce(prev, builder.ConstantR0<float>(0.0f), add,
214 auto test = builder.Gt(builder.ConstantR0<float>(15.5f), sum);
215 condition = builder.Build().ConsumeValueOrDie();
222 ComputationBuilder builder(client_, "body");
223 auto prev = builder.Parameter(0, result_shape, "prev");
224 auto input = builder.ConstantR1<float>({});
225 auto result = builder.Add(input, prev);
226 body = builder.Build().ConsumeValueOrDie();
230 ComputationBuilder builder(client_, "while");
231 auto init = builder.ConstantR1<float>({});
232 auto result = builder.While(condition, body, init);
234 *builder.GetShape(result).ConsumeValueOrDie());
236 ComputeAndCompareR1<float>(&builder, {}, {}, ErrorSpec(0.0001));
252 ComputationBuilder builder(client_, "add");
253 auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
254 auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y");
255 builder.Add(x, y);
256 add = builder.Build().ConsumeValueOrDie();
263 ComputationBuilder builder(client_, "condition");
264 auto prev = builder.Parameter(0, result_shape, "prev");
265 auto sum = builder.Reduce(prev, builder.ConstantR0<float>(0.0f), add,
267 auto test = builder.Gt(builder.ConstantR0<float>(15.5f), sum);
268 condition = builder.Build().ConsumeValueOrDie();
275 ComputationBuilder builder(client_, "body");
276 auto prev = builder.Parameter(0, result_shape, "prev");
277 auto input = builder.ConstantR1<float>(8, 0.125f);
278 auto result = builder.Add(input, prev);
279 body = builder.Build().ConsumeValueOrDie();
283 ComputationBuilder builder(client_, "while");
284 auto init = builder.ConstantR1<float>(8, 0.f);
285 auto result = builder.While(condition, body, init);
287 *builder.GetShape(result).ConsumeValueOrDie());
293 ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
311 ComputationBuilder builder(client_, "add");
312 auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
313 auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y");
314 builder.Add(x, y);
315 add = builder.Build().ConsumeValueOrDie();
322 ComputationBuilder builder(client_, "condition");
323 auto prev = builder.Parameter(0, result_shape, "prev");
324 auto sum = builder.Reduce(prev, builder.ConstantR0<float>(0.0f), add,
326 auto test = builder.Gt(builder.ConstantR0<float>(15.5f), sum);
327 condition = builder.Build().ConsumeValueOrDie();
334 ComputationBuilder builder(client_, "body");
335 auto prev = builder.Parameter(0, result_shape, "prev");
336 auto input = builder.ConstantR1<float>(8, 0.125f);
337 auto result = builder.Add(input, prev);
338 body = builder.Build().ConsumeValueOrDie();
342 ComputationBuilder builder(client_, "while");
343 auto init = builder.ConstantR1<float>(8, 0.f);
344 auto result = builder.While(condition, body, init);
347 *builder.GetShape(result).ConsumeValueOrDie());
348 builder.Tuple({result});
357 ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001));
371 ComputationBuilder builder(client_, "condition");
372 auto prev = builder.Parameter(0, result_shape, "prev");
373 auto iteration = builder.GetTupleElement(prev, 0);
374 builder.Gt(builder.ConstantR0<int32>(N), iteration);
375 condition = builder.Build().ConsumeValueOrDie();
382 ComputationBuilder builder(client_, "body");
383 auto prev = builder.Parameter(0, result_shape, "prev");
384 auto iteration = builder.GetTupleElement(prev, 0);
385 auto w1 = builder.GetTupleElement(prev, 1);
386 auto w2 = builder.GetTupleElement(prev, 2);
387 auto w3 = builder.GetTupleElement(prev, 3);
388 auto result = builder.Tuple(
389 {builder.Add(iteration, builder.ConstantR0<int32>(1)), w3, w1, w2});
390 body = builder.Build().ConsumeValueOrDie();
394 ComputationBuilder builder(client_, "while");
395 auto init = builder.Tuple(
396 {builder.ConstantR0<int32>(0), builder.ConstantR1<float>(3, 1.f),
397 builder.ConstantR1<float>(3, 2.f), builder.ConstantR1<float>(3, 3.f)});
398 auto result = builder.While(condition, body, init);
401 *builder.GetShape(result).ConsumeValueOrDie());
410 ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001));
424 ComputationBuilder builder(client_, "condition");
425 builder.Parameter(0, result_shape, "prev");
426 auto iteration = builder.GetTupleElement(prev, 0);
427 builder.Gt(builder.ConstantR0<int32>(N), iteration);
428 condition = builder.Build().ConsumeValueOrDie();
435 ComputationBuilder builder(client_, "body");
436 auto prev = builder.Parameter(0, result_shape, "prev");
437 auto iteration = builder.GetTupleElement(prev, 0);
438 auto w1 = builder.GetTupleElement(prev, 1);
439 auto w2 = builder.GetTupleElement(prev, 2);
440 auto w3 = builder.GetTupleElement(prev, 3);
441 auto result = builder.Tuple(
442 {builder.Add(iteration, builder.ConstantR0<int32>(1)), w3, w1, w2});
443 body = builder.Build().ConsumeValueOrDie();
447 ComputationBuilder builder(client_, "while");
448 auto init = builder.Tuple(
449 {builder.ConstantR0<int32>(0), builder.ConstantR1<float>(3, 1.f),
450 builder.ConstantR1<float>(3, 2.f), builder.ConstantR1<float>(3, 3.f)});
451 auto xla_while = builder.While(condition, body, init);
453 auto add12 = builder.Add(builder.GetTupleElement(xla_while, 1),
454 builder.GetTupleElement(xla_while, 2));
455 auto result = builder.Add(add12, builder.GetTupleElement(xla_while, 3));
458 *builder.GetShape(result).ConsumeValueOrDie());
460 ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
479 ComputationBuilder builder(client_, "condition");
480 auto prev = builder.Parameter(0, result_shape, "prev");
481 auto iteration = builder.GetTupleElement(prev, 0);
482 builder.Gt(builder.ConstantR0<int32>(5), iteration);
483 condition = builder.Build().ConsumeValueOrDie();
491 ComputationBuilder builder(client_, "body");
492 auto prev = builder.Parameter(0, result_shape, "prev");
493 auto iteration = builder.GetTupleElement(prev, 0);
494 auto weights = builder.GetTupleElement(prev, 1);
495 auto input = builder.ConstantR1<float>(10, 1.f);
496 auto new_weights = builder.Add(weights, input);
497 auto result = builder.Tuple(
498 {builder.Add(iteration, builder.ConstantR0<int32>(1)), new_weights});
499 body = builder.Build().ConsumeValueOrDie();
503 ComputationBuilder builder(client_, "while");
504 auto init = builder.Tuple(
505 {builder.ConstantR0<int32>(0), builder.ConstantR1<float>(10, 0.f)});
506 auto result = builder.While(condition, body, init);
508 *builder.GetShape(result).ConsumeValueOrDie());
516 ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001));
528 ComputationBuilder builder(client_, "condition");
529 auto prev = builder.Parameter(0, result_shape, "prev");
530 auto iteration = builder.GetTupleElement(prev, 0);
531 builder.Gt(builder.ConstantR0<int32>(5), iteration);
532 condition = builder.Build().ConsumeValueOrDie();
539 ComputationBuilder builder(client_, "body");
540 auto prev = builder.Parameter(0, result_shape, "prev");
541 auto iteration = builder.GetTupleElement(prev, 0);
542 auto pred = builder.GetTupleElement(prev, 1);
543 auto new_pred = builder.Or(pred, builder.ConstantR0<bool>(true));
544 auto result = builder.Tuple(
545 {builder.Add(iteration, builder.ConstantR0<int32>(1)), new_pred});
546 body = builder.Build().ConsumeValueOrDie();
550 ComputationBuilder builder(client_, "while");
551 auto init = builder.Tuple({builder.ConstantR0<int32>(0),
552 builder.Ne(builder.ConstantR0<bool>(false),
553 builder.ConstantR0<bool>(true))});
554 auto result = builder.While(condition, body, init);
557 *builder.GetShape(result).ConsumeValueOrDie());
563 ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0));
575 ComputationBuilder builder(client_, "condition");
576 auto prev = builder.Parameter(0, result_shape, "prev");
577 auto iteration = builder.GetTupleElement(prev, 0);
578 builder.Gt(builder.ConstantR0<int32>(5), iteration);
579 condition = builder.Build().ConsumeValueOrDie();
587 ComputationBuilder builder(client_, "body");
588 auto prev = builder.Parameter(0, result_shape, "prev");
589 auto iteration = builder.GetTupleElement(prev, 0);
591 builder.Tuple({builder.Add(iteration, builder.ConstantR0<int32>(1)),
592 builder.ConstantR0<int32>(7)});
593 body = builder.Build().ConsumeValueOrDie();
597 ComputationBuilder builder(client_, "while");
598 auto init = builder.Tuple(
599 {builder.ConstantR0<int32>(0), builder.ConstantR0<int32>(7)});
600 auto result = builder.While(condition, body, init);
603 *builder.GetShape(result).ConsumeValueOrDie());
610 ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001));
637 ComputationBuilder builder(client_, "condition");
638 auto prev = builder.Parameter(0, result_shape, "prev");
639 auto iteration = builder.GetTupleElement(prev, 0);
640 builder.Lt(iteration, builder.ConstantR0<int32>(c1));
641 TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build());
647 ComputationBuilder builder(client_, "condition2");
648 auto prev = builder.Parameter(0, result_shape, "prev");
649 auto iteration = builder.GetTupleElement(prev, 0);
650 builder.Lt(iteration, builder.ConstantR0<int32>(c2));
651 TF_ASSERT_OK_AND_ASSIGN(condition2, builder.Build());
659 ComputationBuilder builder(client_, "body");
660 auto prev = builder.Parameter(0, result_shape, "prev");
661 auto iteration = builder.GetTupleElement(prev, 0);
662 auto weights = builder.GetTupleElement(prev, 1);
663 auto input = builder.ConstantR1<float>(10, 1.f);
664 auto new_weights = builder.Add(weights, input);
665 auto result = builder.Tuple(
666 {builder.Add(iteration, builder.ConstantR0<int32>(1)), new_weights});
667 TF_ASSERT_OK_AND_ASSIGN(body, builder.Build());
672 ComputationBuilder builder(client_, "body");
673 auto prev = builder.Parameter(0, result_shape, "prev");
674 auto iteration = builder.GetTupleElement(prev, 0);
675 auto weights = builder.GetTupleElement(prev, 1);
676 auto input = builder.ConstantR1<float>(10, 1.f);
677 auto new_weights = builder.Add(weights, input);
678 auto result = builder.Tuple(
679 {builder.Add(iteration, builder.ConstantR0<int32>(1)), new_weights});
680 TF_ASSERT_OK_AND_ASSIGN(body2, builder.Build());
684 ComputationBuilder builder(client_, "while");
685 auto init = builder.Tuple(
686 {builder.ConstantR0<int32>(0), builder.ConstantR1<float>(10, 0.f)});
687 auto while1 = builder.While(condition, body, init);
689 auto while2 = builder.While(condition2, body2, while1);
691 auto while_result1 = builder.GetTupleElement(while1, 1);
692 auto while_result2 = builder.GetTupleElement(while2, 1);
695 *builder.GetShape(while_result2).ConsumeValueOrDie());
696 auto result = builder.Add(while_result1, while_result2);
699 *builder.GetShape(result).ConsumeValueOrDie());
702 ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
716 ComputationBuilder builder(client_, "condition");
717 auto prev = builder.Parameter(0, result_shape, "prev");
718 auto iteration = builder.GetTupleElement(prev, 0);
719 builder.Lt(iteration, builder.ConstantR0<int32>(c1));
720 TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build());
726 ComputationBuilder builder(client_, "condition2");
727 auto prev = builder.Parameter(0, result_shape, "prev");
728 auto iteration = builder.GetTupleElement(prev, 0);
729 builder.Lt(iteration, builder.ConstantR0<int32>(c2));
730 TF_ASSERT_OK_AND_ASSIGN(condition2, builder.Build());
738 ComputationBuilder builder(client_, "body");
739 auto prev = builder.Parameter(0, result_shape, "prev");
740 auto iteration = builder.GetTupleElement(prev, 0);
741 auto weights = builder.GetTupleElement(prev, 1);
742 auto input = builder.ConstantR1<float>(10, 1.f);
743 auto new_weights = builder.Add(weights, input);
744 auto result = builder.Tuple(
745 {builder.Add(iteration, builder.ConstantR0<int32>(1)), new_weights});
746 TF_ASSERT_OK_AND_ASSIGN(body, builder.Build());
750 ComputationBuilder builder(client_, "while");
751 auto init = builder.Tuple(
752 {builder.ConstantR0<int32>(0), builder.ConstantR1<float>(10, 0.f)});
753 auto while1 = builder.While(condition, body, init);
755 auto while2 = builder.While(condition2, body, while1);
757 auto while_result1 = builder.GetTupleElement(while1, 1);
758 auto while_result2 = builder.GetTupleElement(while2, 1);
761 *builder.GetShape(while_result2).ConsumeValueOrDie());
762 auto result = builder.Add(while_result1, while_result2);
765 *builder.GetShape(result).ConsumeValueOrDie());
768 ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
783 ComputationBuilder builder(client_, "condition");
784 auto prev = builder.Parameter(0, result_shape, "prev");
785 auto iteration = builder.GetTupleElement(prev, 0);
786 builder.Lt(iteration, builder.ConstantR0<int32>(c1));
787 TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build());
793 ComputationBuilder builder(client_, "condition2");
794 auto prev = builder.Parameter(0, result_shape, "prev");
795 auto iteration = builder.GetTupleElement(prev, 0);
796 builder.Lt(iteration, builder.ConstantR0<int32>(c2));
797 TF_ASSERT_OK_AND_ASSIGN(condition2, builder.Build());
805 ComputationBuilder builder(client_, "body");
806 auto prev = builder.Parameter(0, result_shape, "prev");
807 auto iteration = builder.GetTupleElement(prev, 0);
808 auto weights = builder.GetTupleElement(prev, 1);
809 auto input = builder.ConstantR1<float>(10, 1.f);
810 auto new_weights = builder.Add(weights, input);
811 auto result = builder.Tuple(
812 {builder.Add(iteration, builder.ConstantR0<int32>(1)), new_weights});
813 TF_ASSERT_OK_AND_ASSIGN(body, builder.Build());
817 ComputationBuilder builder(client_, "while");
818 builder.Tuple(
819 {builder.ConstantR0<int32>(0), builder.ConstantR1<float>(10, 0.f)});
820 auto while1 = builder.While(condition, body, init);
821 auto while2 = builder.While(condition2, body, init);
823 auto while_result1 = builder.GetTupleElement(while1, 1);
824 auto while_result2 = builder.GetTupleElement(while2, 1);
827 *builder.GetShape(while_result2).ConsumeValueOrDie());
828 auto result = builder.Add(while_result1, while_result2);
831 *builder.GetShape(result).ConsumeValueOrDie());
834 ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
849 ComputationBuilder builder(client_, "condition");
850 auto prev = builder.Parameter(0, result_shape, "prev");
851 auto iteration = builder.GetTupleElement(prev, 0);
852 builder.Gt(builder.ConstantR0<int32>(5), iteration);
853 condition = builder.Build().ConsumeValueOrDie();
861 ComputationBuilder builder(client_, "body");
862 auto prev = builder.Parameter(0, result_shape, "prev");
864 auto iteration = builder.GetTupleElement(prev, 0);
865 auto out0 = builder.Add(iteration, builder.ConstantR0<int32>(1));
867 auto input = builder.GetTupleElement(prev, 1);
869 auto update = builder.ConvertElementType(builder.Broadcast(out0, {2}), F32);
871 auto starts = builder.Reshape(
872 builder.Mul(iteration, builder.ConstantR0<int32>(2)), {1});
874 auto out1 = builder.DynamicUpdateSlice(input, update, starts);
876 auto result = builder.Tuple({out0, out1});
877 body = builder.Build().ConsumeValueOrDie();
881 ComputationBuilder builder(client_, "while");
882 auto init = builder.Tuple(
883 {builder.ConstantR0<int32>(0), builder.ConstantR1<float>(10, 0.f)});
884 auto result = builder.While(condition, body, init);
887 *builder.GetShape(result).ConsumeValueOrDie());
895 ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001));
918 ComputationBuilder builder(client_, TestName());
919 auto prev = builder.Reshape(
920 builder.Slice(builder.Parameter(0, v6s32, "prev"), {0}, {1}, {1}), {0},
922 builder.Gt(builder.ConstantR0<int32>(count), prev);
923 return builder.Build().ConsumeValueOrDie();
929 ComputationBuilder builder(client_, "body");
930 auto prev = builder.Parameter(0, v6s32, "prev");
931 auto inc = builder.ConcatInDim(
932 {builder.ConstantR1<int32>({1}),
933 builder.RngUniform(builder.ConstantR0<int32>(0),
934 builder.ConstantR0<int32>(100),
937 auto result = builder.Add(inc, prev);
938 body = builder.Build().ConsumeValueOrDie();
943 ComputationBuilder builder(client_, TestName());
944 auto init = builder.ConstantR1<int32>({0, 0, 0, 0, 0, 0});
945 auto result = builder.While(build_condition(count), body, init);
946 auto shape = builder.GetShape(result).ConsumeValueOrDie();
947 return builder.Build();
1112 ComputationBuilder builder(client_, "inner_condition");
1113 auto params = builder.Parameter(0, inner_result_shape, "prev");
1114 auto i = builder.GetTupleElement(params, 0);
1115 builder.Lt(i, builder.ConstantR0<int32>(7));
1116 inner_condition = builder.Build().ConsumeValueOrDie();
1123 ComputationBuilder builder(client_, "outer_condition");
1124 auto prev = builder.Parameter(0, outer_result_shape, "prev");
1125 builder.Lt(prev, builder.ConstantR0<int32>(30));
1126 outer_condition = builder.Build().ConsumeValueOrDie();
1133 ComputationBuilder builder(client_, "inner_body");
1134 auto params = builder.Parameter(0, inner_result_shape, "prev");
1135 auto i = builder.GetTupleElement(params, 0);
1136 auto result = builder.GetTupleElement(params, 1);
1137 i = builder.Add(builder.ConstantR0<int32>(1), i);
1138 result = builder.Add(builder.ConstantR0<int32>(2), result);
1139 auto output = builder.Tuple({i, result});
1140 inner_body = builder.Build().ConsumeValueOrDie();
1146 ComputationBuilder builder(client_, "outer_body");
1147 auto prev = builder.Parameter(0, outer_result_shape, "prev");
1148 auto init = builder.Tuple({builder.ConstantR0<int32>(0), prev});
1149 auto result = builder.While(inner_condition, inner_body, init);
1150 auto output = builder.GetTupleElement(result, 1);
1151 outer_body = builder.Build().ConsumeValueOrDie();
1155 ComputationBuilder builder(client_, TestName());
1156 auto init = builder.ConstantR0<int32>(0);
1157 auto result = builder.While(outer_condition, outer_body, init);
1158 auto shape = builder.GetShape(result).ConsumeValueOrDie();
1160 ComputeAndCompareR0<int32>(&builder, 42, {});
1175 ComputationBuilder builder(client_, "condition_callee");
1176 auto prev = builder.Parameter(0, result_shape, "prev");
1177 builder.Tuple({builder.Gt(builder.ConstantR0<int32>(5), prev)});
1179 condition_callee = builder.Build().ConsumeValueOrDie();
1184 ComputationBuilder builder(client_, "condition");
1185 auto prev = builder.Parameter(0, result_shape, "prev");
1186 auto result = builder.Call(condition_callee, {prev});
1187 builder.GetTupleElement(result, 0);
1188 condition = builder.Build().ConsumeValueOrDie();
1194 ComputationBuilder builder(client_, "body");
1195 auto prev = builder.Parameter(0, result_shape, "prev");
1196 auto input = builder.ConstantR0<int32>(1);
1197 auto result = builder.Add(input, prev);
1198 body = builder.Build().ConsumeValueOrDie();
1202 ComputationBuilder builder(client_, TestName());
1203 auto init = builder.ConstantR0<int32>(0);
1204 auto result = builder.While(condition, body, init);
1205 auto shape = builder.GetShape(result).ConsumeValueOrDie();
1207 ComputeAndCompareR0<int32>(&builder, 5, {});
1219 ComputationBuilder builder(client_, "condition");
1220 auto state = builder.Parameter(0, while_shape, "state");
1221 builder.Gt(builder.ConstantR0<int32>(5), builder.GetTupleElement(state, 0));
1222 TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build());
1227 ComputationBuilder builder(client_, "body");
1228 auto state = builder.Parameter(0, while_shape, "state");
1229 auto indvar = builder.GetTupleElement(state, 0);
1230 auto input_0 = builder.GetTupleElement(state, 1);
1231 auto input_1 = builder.GetTupleElement(state, 2);
1232 auto output = builder.Tanh(builder.Dot(input_0, input_1));
1233 auto indvar_next = builder.Add(indvar, builder.ConstantR0<int32>(1));
1234 auto tuple_result = builder.Tuple({indvar_next, input_0, input_1, output});
1235 builder.Build());
1238 ComputationBuilder builder(client_, TestName());
1239 auto matrix_input = builder.Parameter(0, matrix_shape, "matrix");
1240 auto init = builder.Tuple(
1241 {builder.ConstantR0<int32>(0), matrix_input, matrix_input, matrix_input});
1242 auto while_instruction = builder.While(condition, body, init);
1243 builder.GetTupleElement(while_instruction, 3);
1250 &builder, {{-0.76159416, -0.96402758}, {0.76159416, 0.96402758}},
1273 ComputationBuilder builder(client, "condition");
1274 auto prev = builder.Parameter(0, loop_state_shape, "prev");
1275 auto iteration = builder.GetTupleElement(prev, 0);
1276 builder.Lt(iteration, builder.ConstantR0<int32>(loop_limit));
1277 condition = builder.Build().ConsumeValueOrDie();
1283 ComputationBuilder builder(client, "body");
1284 auto prev = builder.Parameter(0, loop_state_shape, "prev");
1286 auto iteration = builder.GetTupleElement(prev, 0);
1287 auto out0 = builder.Add(iteration, builder.ConstantR0<int32>(1));
1289 auto input = builder.GetTupleElement(prev, 1);
1291 auto one = builder.ConstantR0<float>(1.0);
1292 auto update = builder.Broadcast(one, {1, 1024, 1024});
1294 auto starts = builder.ConstantR1<int32>({0, 0, 0});
1296 auto out1 = builder.DynamicUpdateSlice(input, update, starts);
1297 auto result = builder.Tuple({out0, out1});
1298 body = builder.Build().ConsumeValueOrDie();
1302 ComputationBuilder builder(client, "while");
1303 auto zero = builder.ConstantR0<float>(0.0);
1304 auto input = builder.Broadcast(zero, {seq_len, 1024, 1024});
1305 auto init = builder.Tuple({builder.ConstantR0<int32>(0), input});
1306 builder.While(condition, body, init);
1307 auto computation = builder.Build().ConsumeValueOrDie();
