Coverage for mlprodict/onnxrt/ops_cpu/op_eyelike.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 

9from ._op_helper import proto2dtype, dtype_name 

10 

11 

12class EyeLike(OpRun): 

13 

14 atts = {'k': 0, 'dtype': 1} 

15 

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

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

18 expected_attributes=EyeLike.atts, 

19 **options) 

20 self.dtype_ = proto2dtype(self.dtype) 

21 

22 def _run(self, data, *args, attributes=None, verbose=0, fLOG=None): # pylint: disable=W0221 

23 shape = data.shape 

24 if len(shape) == 1: 

25 sh = (shape[0], shape[0]) 

26 elif len(shape) == 2: 

27 sh = shape 

28 else: 

29 raise RuntimeError( # pragma: no cover 

30 f"EyeLike only accept 1D or 2D tensors not {shape!r}.") 

31 return (numpy.eye(*sh, k=self.k, dtype=self.dtype_), ) 

32 

33 def to_python(self, inputs): 

34 return ( 

35 "import numpy", 

36 "return numpy.eye(*(%s.shape), k=%d, dtype=numpy.%s)" % ( 

37 inputs[0], self.k, dtype_name(self.dtype_)))