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 

10from ..shape_object import ShapeObject 

11 

12 

13class EyeLike(OpRun): 

14 

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

16 

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

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

19 expected_attributes=EyeLike.atts, 

20 **options) 

21 self.dtype_ = proto2dtype(self.dtype) 

22 

23 def _run(self, data, *args): # pylint: disable=W0221 

24 shape = data.shape 

25 if len(shape) == 1: 

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

27 elif len(shape) == 2: 

28 sh = shape 

29 else: 

30 raise RuntimeError( # pragma: no cover 

31 "EyeLike only accept 1D or 2D tensors not %r." % (shape, )) 

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

33 

34 def _infer_shapes(self, data): # pylint: disable=W0221 

35 return (ShapeObject(None, dtype=self.dtype_), ) 

36 

37 def _infer_types(self, data): # pylint: disable=W0221 

38 return (self.dtype_, ) 

39 

40 def _infer_sizes(self, *args): # pylint: disable=W0221 

41 res = self.run(*args) 

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

43 

44 def to_python(self, inputs): 

45 return ( 

46 "import numpy", 

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

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