Ctrl+K

nir.ir.utils

nir.ir.utils#

Module Contents#

Functions#

parse_shape_argument(x, key)

Parse the shape argument of a NIR node.

calculate_conv_output(→ Sequence[int])

Calculates the output for a single dimension of a convolutional layer.

calc_flatten_output(input_shape, start_dim, end_dim)

ensure_str(→ str)

nir.ir.utils.parse_shape_argument(x: nir.ir.typing.Types, key: str)#

Parse the shape argument of a NIR node.

Parameters

  • x (nir.ir.typing.Types) –

  • key (str) –

nir.ir.utils.calculate_conv_output(input_shape: Union[int, Sequence[int]], padding: Union[int, str, Sequence[int]], dilation: Union[int, Sequence[int]], kernel_size: Union[int, Sequence[int]], stride: Union[int, Sequence[int]]) → Sequence[int]#

Calculates the output for a single dimension of a convolutional layer. https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html#torch.nn.Conv1d

Parameters

  • input_shape (int | Sequence[int]) – input shape, either int or (int, int)

  • padding (int | Sequence[int]) – padding

  • dilation (int | Sequence[int]) – dilation

  • kernel_size (int | Sequence[int]) – kernel size

  • stride (int | Sequence[int]) – stride

Returns

output shape

Return type

Sequence[int]

nir.ir.utils.calc_flatten_output(input_shape: Sequence[int], start_dim: int, end_dim: int)#
Parameters

  • input_shape (Sequence[int]) –

  • start_dim (int) –

  • end_dim (int) –

nir.ir.utils.ensure_str(a: Union[str, bytes]) → str#
Parameters

a (Union[str, bytes]) –

Return type

str

Contents