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, shape, *args): # pylint: disable=W0221 

24 return (numpy.eye(*shape, k=self.k, dtype=self.dtype_), ) 

25 

26 def _infer_shapes(self, shape): # pylint: disable=W0221 

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

28 

29 def _infer_types(self, shape): # pylint: disable=W0221 

30 return (self.dtype_, ) 

31 

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

33 res = self.run(*args) 

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

35 

36 def to_python(self, inputs): 

37 return ( 

38 "import numpy", 

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

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