()
| 41 | |
| 42 | |
| 43 | def test_block_builder(): |
| 44 | m = tirx.Var("m", "int64") |
| 45 | n = tirx.Var("n", "int64") |
| 46 | x = rx.Var("x", rx.TensorStructInfo([m, n], "float16")) |
| 47 | y = rx.Var("y", rx.TensorStructInfo([n], "float16")) |
| 48 | bb = rx.BlockBuilder() |
| 49 | |
| 50 | bb._begin_binding_block() |
| 51 | gv0 = bb.emit(rx.op.add(x, y)) |
| 52 | bb._begin_dataflow_block() |
| 53 | lv0 = bb.emit(rx.op.multiply(gv0, y)) |
| 54 | gv1 = bb.emit_output(rx.op.multiply(lv0, lv0)) |
| 55 | b0 = bb._end_block() |
| 56 | bb._begin_dataflow_block() |
| 57 | lv1 = bb.emit(rx.op.multiply(gv0, y)) |
| 58 | gv2 = bb.emit_output(rx.op.multiply(lv1, lv1)) |
| 59 | b1 = bb._end_block() |
| 60 | gv3 = bb.emit(rx.op.add(x, y)) |
| 61 | b2 = bb._end_block() |
| 62 | |
| 63 | assert isinstance(b0, rx.DataflowBlock) |
| 64 | assert isinstance(b1, rx.DataflowBlock) |
| 65 | assert not isinstance(b2, rx.DataflowBlock) |
| 66 | |
| 67 | |
| 68 | def test_emit_with_name(): |
nothing calls this directly
no test coverage detected
searching dependent graphs…