(is_varnode)
| 39 | |
| 40 | @pytest.mark.parametrize("is_varnode", [False, True]) |
| 41 | def test_diag(is_varnode): |
| 42 | if is_varnode: |
| 43 | network = Network() |
| 44 | else: |
| 45 | network = None |
| 46 | |
| 47 | shapes = [(10, 10), (6, 9), (8, 7), (8,)] |
| 48 | cases = [] |
| 49 | for shp in shapes: |
| 50 | cases.append({"input": [np.random.random(shp).astype("float32")]}) |
| 51 | |
| 52 | for axis in range(-2, 3): |
| 53 | |
| 54 | def run(data): |
| 55 | return F.diag(data, k=axis) |
| 56 | |
| 57 | opr_test(cases, run, ref_fn=lambda x: np.diag(x, axis), network=network) |
| 58 | |
| 59 | |
| 60 | def test_full(): |