(shape, dim, in_layout, expected_out_layout)
| 294 | |
| 295 | def test_tensor_cpu_squeeze(): |
| 296 | def check_squeeze(shape, dim, in_layout, expected_out_layout): |
| 297 | arr = np.random.rand(*shape) |
| 298 | t = TensorCPU(arr, in_layout) |
| 299 | is_squeezed = t.squeeze(dim) |
| 300 | should_squeeze = len(expected_out_layout) < len(in_layout) |
| 301 | arr_squeeze = arr.squeeze(dim) |
| 302 | t_shape = tuple(t.shape()) |
| 303 | assert t_shape == arr_squeeze.shape, f"{t_shape} != {arr_squeeze.shape}" |
| 304 | assert t.layout() == expected_out_layout, f"{t.layout()} != {expected_out_layout}" |
| 305 | assert layout_compatible( |
| 306 | t.get_property("layout"), expected_out_layout |
| 307 | ), f'{t.get_property("layout")} doesn\'t match {expected_out_layout}' |
| 308 | assert np.allclose(arr_squeeze, np.array(t)) |
| 309 | assert is_squeezed == should_squeeze, f"{is_squeezed} != {should_squeeze}" |
| 310 | |
| 311 | for dim, shape, in_layout, expected_out_layout in [ |
| 312 | (None, (3, 5, 6), "ABC", "ABC"), |
nothing calls this directly
no test coverage detected