Appends dimensions to the end of a tensor until it has target_dims dimensions.
(x, target_dims)
| 42 | |
| 43 | |
| 44 | def append_dims(x, target_dims): |
| 45 | """Appends dimensions to the end of a tensor until it has target_dims dimensions.""" |
| 46 | dims_to_append = target_dims - x.ndim |
| 47 | if dims_to_append < 0: |
| 48 | raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less') |
| 49 | return x[(...,) + (None,) * dims_to_append] |
| 50 | |
| 51 | |
| 52 | def n_params(module): |
nothing calls this directly
no outgoing calls
no test coverage detected