Coverage for mlprodict/onnxrt/ops_cpu/op_layer_normalization.py: 100%

1# -*- encoding: utf-8 -*- 

2# pylint: disable=E0203,E1101,C0111 

3""" 

4@file 

5@brief Runtime operator. 

6""" 

7import numpy 

8from ._op import OpRun 

9 

10 

11def _layer_normalization(X, W, B, axis=-1, epsilon=1e-5): # type: ignore 

12 # Inspired from: https://github.com/onnx/onnx/blob/main/onnx/backend/ 

13 # test/case/node/layernormalization.py#L12 

14 X_shape = X.shape 

15 X_rank = len(X_shape) 

16 if axis < 0: 

17 # If axis = -1 and rank of X is 4, 

18 # the axis is changed to -1 + 4 = 3, 

19 # which means the last axis. 

20 axis = axis + X_rank 

21 unsqueezed_rank = X_rank - axis 

22 reduction_shape = X_shape[0:axis] + (1,) * unsqueezed_rank 

23 

24 # Parameter used to convert N-D tensor layer 

25 # normalization to equivalent 2-D matirx operations. 

26 row_number = 1 

27 col_number = 1 

28 for i in range(X_rank): 

29 if i < axis: 

30 row_number *= X_shape[i] 

31 else: 

32 col_number *= X_shape[i] 

33 

34 # After reshaping input tensor X into a matrix, 

35 # layer normalization is equivalent to conducting 

36 # standardization on each column vector (s.t. each 

37 # column has zero mean and unit variance). 

38 x_mat = numpy.reshape(X, (row_number, col_number)) 

39 # This computes mean for every x_mat's column. 

40 x_mean = numpy.sum(x_mat, axis=1, keepdims=True) / col_number 

41 x_diff = x_mat - x_mean 

42 x_squared_diff = x_diff * x_diff 

43 # This computes variance for every x_mat's column. 

44 variance = numpy.sum(x_squared_diff, axis=1, keepdims=True) / col_number 

45 variance_eps = variance + epsilon 

46 std_dev = numpy.sqrt(variance_eps) 

47 inv_std_dev = numpy.reciprocal(std_dev) 

48 # Standardization step. y_mat is zero-mean and unit-variance. 

49 y_mat = x_diff * inv_std_dev 

50 # Apply affine transform on normalization outcome. 

51 # W is linear coefficient while B is bias. 

52 Y = numpy.reshape(y_mat, X_shape) * W 

53 if B is not None: 

54 Y = Y + B 

55 # Matrix-level operations' outputs should be reshaped 

56 # to compensate the initial tensor-to-matrix reshape. 

57 X_mean = numpy.reshape(x_mean, reduction_shape) 

58 X_inv_std_dev = numpy.reshape(inv_std_dev, reduction_shape) 

59 

60 return (Y.astype(X.dtype), 

61 X_mean.astype(X.dtype), 

62 X_inv_std_dev.astype(X.dtype)) 

63 

64 

65class LayerNormalization(OpRun): 

66 

67 atts = {'axis': -1, 

68 'epsilon': 9.999999747378752e-06, 

69 'stash_type': 1} 

70 

71 def __init__(self, onnx_node, desc=None, **options): 

72 OpRun.__init__(self, onnx_node, desc=desc, 

73 expected_attributes=LayerNormalization.atts, 

74 **options) 

75 

76 def _run(self, X, Scale, B=None, attributes=None, verbose=0, fLOG=None): # pylint: disable=W0221 

77 res = _layer_normalization( 

78 X, Scale, B, axis=self.axis, epsilon=self.epsilon) 

79 return res