(tensor, ref_shape)
| 95 | # Performs symbolic assertion when used in torch.jit.trace(). |
| 96 | |
| 97 | def assert_shape(tensor, ref_shape): |
| 98 | if tensor.ndim != len(ref_shape): |
| 99 | raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}') |
| 100 | for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)): |
| 101 | if ref_size is None: |
| 102 | pass |
| 103 | elif isinstance(ref_size, torch.Tensor): |
| 104 | with suppress_tracer_warnings(): # as_tensor results are registered as constants |
| 105 | symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}') |
| 106 | elif isinstance(size, torch.Tensor): |
| 107 | with suppress_tracer_warnings(): # as_tensor results are registered as constants |
| 108 | symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}') |
| 109 | elif size != ref_size: |
| 110 | raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}') |
| 111 | |
| 112 | #---------------------------------------------------------------------------- |
| 113 | # Function decorator that calls torch.autograd.profiler.record_function(). |
nothing calls this directly
no test coverage detected