Coverage for mlprodict/onnxrt/ops_cpu/op_trilu.py: 96%

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 

11class Trilu(OpRun): 

12 

13 atts = {'upper': 1} 

14 

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

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

17 expected_attributes=Trilu.atts, 

18 **options) 

19 if self.upper not in (0, 1): 

20 raise ValueError("upper must be 0 or 1 not %r." % (self.upper, )) 

21 

22 def _run(self, *inputs): # pylint: disable=W0221 

23 x = inputs[0] 

24 k = 0 if len(inputs) == 1 else int(inputs[1]) 

25 if self.upper: 

26 return (numpy.triu(x, k), ) 

27 return (numpy.tril(x, k), ) 

28 

29 def to_python(self, inputs): 

30 name = "triu" if self.upper else "tril" 

31 return ( 

32 "import numpy", 

33 "return numpy.%s(%s, int(%s))" % ( 

34 name, inputs[0], 0 if len(inputs) == 1 else inputs[1])) 

35 

36 def _infer_shapes(self, *inputs): # pylint: disable=W0221 

37 return (inputs[0], ) 

38 

39 def _infer_types(self, *inputs): # pylint: disable=W0221 

40 return (inputs[0], ) 

41 

42 def _infer_sizes(self, *args, **kwargs): 

43 res = self.run(*args, **kwargs) 

44 return (dict(temp=0), ) + res