Coverage for mlprodict/onnxrt/ops_cpu/op_where.py: 100%

10 statements  

« prev     ^ index     » next       coverage.py v7.1.0, created at 2023-02-04 02:28 +0100

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 Where(OpRun): 

12 

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

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

15 **options) 

16 

17 def _run(self, condition, x, y, attributes=None, verbose=0, fLOG=None): # pylint: disable=W0221 

18 if x.dtype != y.dtype and x.dtype not in (numpy.object_, ): 

19 raise RuntimeError( # pragma: no cover 

20 f"x and y should share the same dtype {x.dtype} != {y.dtype}") 

21 if x.shape != y.shape and x.shape != (1, ) and y.shape != (1, ): 

22 raise RuntimeError( # pragma: no cover 

23 f"x and y should share the same shape {x.shape} != {y.shape}") 

24 return (numpy.where(condition, x, y).astype(x.dtype), )