concrete.ml.onnx.onnx_impl_utils.md

module concrete.ml.onnx.onnx_impl_utils

Utility functions for onnx operator implementations.

function numpy_onnx_pad

numpy_onnx_pad(
    x: ndarray,
    pads: Tuple[int, ],
    pad_value: Union[float, int, ndarray] = 0,
    int_only: bool = False
) → ndarray

Pad a tensor according to ONNX spec, using an optional custom pad value.

Args:

• x (numpy.ndarray): input tensor to pad

• pads (List[int]): padding values according to ONNX spec

• pad_value (Optional[Union[float, int]]): value used to fill in padding, default 0

• int_only (bool): set to True to generate integer only code with Concrete

Returns:

• res (numpy.ndarray): the input tensor with padding applied

function compute_conv_output_dims

compute_conv_output_dims(
    input_shape: Tuple[int, ],
    kernel_shape: Tuple[int, ],
    pads: Tuple[int, ],
    strides: Tuple[int, ],
    ceil_mode: int
) → Tuple[int, ]

Compute the output shape of a pool or conv operation.

See https://pytorch.org/docs/stable/generated/torch.nn.AvgPool2d.html for details on the computation of the output shape.

Args:

• input_shape (Tuple[int, ...]): shape of the input to be padded as N x C x H x W

• kernel_shape (Tuple[int, ...]): shape of the conv or pool kernel, as Kh x Kw (or n-d)

• pads (Tuple[int, ...]): padding values following ONNX spec: dim1_start, dim2_start, .. dimN_start, dim1_end, dim2_end, ... dimN_end where in the 2-d case dim1 is H, dim2 is W

• strides (Tuple[int, ...]): strides for each dimension

• ceil_mode (int): set to 1 to use the ceil function to compute the output shape, as described in the PyTorch doc

Returns:

• res (Tuple[int, ...]): shape of the output of a conv or pool operator with given parameters

function compute_onnx_pool_padding

compute_onnx_pool_padding(
    input_shape: Tuple[int, ],
    kernel_shape: Tuple[int, ],
    pads: Tuple[int, ],
    strides: Tuple[int, ],
    ceil_mode: int
) → Tuple[int, ]

Compute any additional padding needed to compute pooling layers.

The ONNX standard uses ceil_mode=1 to match TensorFlow style pooling output computation. In this setting, the kernel can be placed at a valid position even though it contains values outside of the input shape including padding. The ceil_mode parameter controls whether this mode is enabled. If the mode is not enabled, the output shape follows PyTorch rules.

Args:

• input_shape (Tuple[int, ...]): shape of the input to be padded as N x C x H x W

• kernel_shape (Tuple[int, ...]): shape of the conv or pool kernel, as Kh x Kw (or n-d)

• pads (Tuple[int, ...]): padding values following ONNX spec: dim1_start, dim2_start, .. dimN_start, dim1_end, dim2_end, ... dimN_end where in the 2-d case dim1 is H, dim2 is W

• strides (Tuple[int, ...]): strides for each dimension

• ceil_mode (int): set to 1 to use the ceil function to compute the output shape, as described in the PyTorch doc

Returns:

• res (Tuple[int, ...]): shape of the output of a conv or pool operator with given parameters

function onnx_avgpool_compute_norm_const

onnx_avgpool_compute_norm_const(
    input_shape: Tuple[int, ],
    kernel_shape: Tuple[int, ],
    pads: Tuple[int, ],
    strides: Tuple[int, ],
    ceil_mode: int
) → Union[ndarray, float]

Compute the average pooling normalization constant.

This constant can be a tensor of the same shape as the input or a scalar.

Args:

• input_shape (Tuple[int, ...]): shape of the input to be padded as N x C x H x W

• kernel_shape (Tuple[int, ...]): shape of the conv or pool kernel, as Kh x Kw (or n-d)

• pads (Tuple[int, ...]): padding values following ONNX spec: dim1_start, dim2_start, .. dimN_start, dim1_end, dim2_end, ... dimN_end where in the 2-d case dim1 is H, dim2 is W

• strides (Tuple[int, ...]): strides for each dimension

• ceil_mode (int): set to 1 to use the ceil function to compute the output shape, as described in the PyTorch doc

Returns:

• res (float): tensor or scalar, corresponding to normalization factors to apply for the average pool computation for each valid kernel position