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

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(f"upper must be 0 or 1 not {self.upper!r}.") 

21 

22 def _run(self, *inputs, attributes=None, verbose=0, fLOG=None): # 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]))