Creates num_decodes as second dimension in non-scalar array x and tiles into it.
(x: Tensor, num_decodes: int)
| 42 | |
| 43 | |
| 44 | def add_decoding_dim(x: Tensor, num_decodes: int) -> Tensor: |
| 45 | """Creates num_decodes as second dimension in non-scalar array x and tiles into it.""" |
| 46 | if x.ndim > 0: |
| 47 | x = jnp.expand_dims(x, axis=1) |
| 48 | tile_dims = [1] * x.ndim |
| 49 | tile_dims[1] = num_decodes |
| 50 | return jnp.tile(x, tile_dims) |
| 51 | |
| 52 | # Scalar values get converted to tensor on jit. |
| 53 | # However, we do not want to handle num_decodes dimension for scalars. |
| 54 | return x |
| 55 | |
| 56 | |
| 57 | def flatten_decoding_dim(x: Tensor) -> Tensor: |
no outgoing calls
no test coverage detected