(tensor, ref_shape)
| 78 | # Performs symbolic assertion when used in torch.jit.trace(). |
| 79 | |
| 80 | def assert_shape(tensor, ref_shape): |
| 81 | err_suffix = f' for tensor of size {list(tensor.shape)}' |
| 82 | if tensor.ndim != len(ref_shape): |
| 83 | raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}{err_suffix}') |
| 84 | for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)): |
| 85 | if ref_size is None: |
| 86 | pass |
| 87 | elif isinstance(ref_size, torch.Tensor): |
| 88 | with suppress_tracer_warnings(): # as_tensor results are registered as constants |
| 89 | symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}{err_suffix}') |
| 90 | elif isinstance(size, torch.Tensor): |
| 91 | with suppress_tracer_warnings(): # as_tensor results are registered as constants |
| 92 | symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}{err_suffix}') |
| 93 | elif size != ref_size: |
| 94 | raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}{err_suffix}') |
| 95 | |
| 96 | #---------------------------------------------------------------------------- |
| 97 | # Function decorator that calls torch.autograd.profiler.record_function(). |
nothing calls this directly
no test coverage detected