WIP: median support for OpenVINO back-end

Author: alungriffith

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -38,7 +38,6 @@ NumpyDtypeTest::test_logspace
 NumpyDtypeTest::test_matmul_
 NumpyDtypeTest::test_max
 NumpyDtypeTest::test_mean
-NumpyDtypeTest::test_median
 NumpyDtypeTest::test_meshgrid
 NumpyDtypeTest::test_minimum_python_types
 NumpyDtypeTest::test_multiply
@@ -95,7 +94,6 @@ NumpyOneInputOpsCorrectnessTest::test_isinf
 NumpyOneInputOpsCorrectnessTest::test_logaddexp
 NumpyOneInputOpsCorrectnessTest::test_max
 NumpyOneInputOpsCorrectnessTest::test_mean
-NumpyOneInputOpsCorrectnessTest::test_median
 NumpyOneInputOpsCorrectnessTest::test_meshgrid
 NumpyOneInputOpsCorrectnessTest::test_pad_float16_constant_2
 NumpyOneInputOpsCorrectnessTest::test_pad_float32_constant_2

keras/src/backend/openvino/numpy.py

@@ -1046,8 +1046,73 @@ def maximum(x1, x2):
 
 
 def median(x, axis=None, keepdims=False):
-    raise NotImplementedError("`median` is not supported with openvino backend")
-
+    x = get_ov_output(x)
+    x_shape_original = ov_opset.shape_of(x).output(0)
+    
+    if axis is None:
+        flatten_shape = ov_opset.constant([-1], Type.i32).output(0)
+        x = ov_opset.reshape(x, flatten_shape, False).output(0)
+        axis = 0
+        flattened = True
+        int_axis = False
+        x_shape = ov_opset.shape_of(x).output(0)
+        k_value = ov_opset.convert(x_shape, Type.i32).output(0)
+    elif isinstance(axis, int):
+        flattened = False
+        int_axis = True
+        ov_axis = ov_opset.constant(axis, Type.i32).output(0)
+        x_shape = ov_opset.shape_of(x).output(0)
+        k_value = ov_opset.convert(ov_opset.gather(x_shape, ov_axis, ov_opset.constant([0], Type.i32).output(0)).output(0), Type.i32).output(0)
+    else:
+        # axis = (2, 1)
+        flattened = False
+        int_axis = False
+        ov_axis = ov_opset.constant(axis, Type.i32).output(0)   # (2, 1)
+        x_rank = ov_opset.shape_of(x_shape_original).output(0)  # 4
+        axis_range = ov_opset.range(ov_opset.constant([0], Type.i32).output(0), x_rank, ov_opset.constant([1], Type.i32).output(0)).output(0)
+        axis_compare = ov_opset.equal(ov_opset.unsqueeze(ov_axis, 1).output(0), ov_opset.unsqueeze(axis_range, 0).output(0)).output(0)
+        mask_remove = ov_opset.reduce_logical_or(axis_compare, ov_opset.constant([0], Type.i32).output(0)).output(0)
+        mask_keep = ov_opset.logical_not(mask_remove).output(0)
+        nz = ov_opset.non_zero(mask_keep, "i32").output(0)
+        indices_keep = ov_opset.squeeze(nz, [0]).output(0)
+        axis_range = ov_opset.gather(axis_range, indices_keep, ov_opset.constant([0], Type.i32).output(0)).output(0) # (0, 3)
+        axis_range = ov_opset.concat([axis_range, ov_axis], ov_opset.constant([0], Type.i32).output(0)).output(0) # (0, 3, 2, 1)
+        x = ov_opset.transpose(x, axis_range).output(0)   # x = (d0, d3, d2, d1)
+        
+        flat_rank = ov_opset.subtract(x_rank, ov_opset.constant([1], Type.i32)).output(0)
+        flatten_shape = ov_opset.constant([0], shape=flat_rank, type_info=Type.i32).output(0)
+        flatten_shape = ov_opset.scatter_elements_update(flatten_shape, ov_opset.constant([-1], Type.i32).output(0), [-1], [0], "sum")
+
+        x = ov_opset.reshape(x, flatten_shape, True).output(0)  # x = (d0, d3, d2*d1)
+        axis = -1
+        x_shape = ov_opset.shape_of(x).output(0)
+        k_value = ov_opset.gather(x_shape, ov_opset.constant([-1], Type.i32).output(0), ov_opset.constant([0], Type.i32).output(0)).output(0)
+        k_value = ov_opset.convert(k_value, Type.i32).output(0)
+        
+    x_sorted = ov_opset.topk(x, k_value, axis, 'min', 'value', stable=True).output(0)
+    half_index = ov_opset.divide(k_value, ov_opset.constant([2], Type.i32)).output(0)
+    x_mod = ov_opset.mod(k_value, ov_opset.constant([2], Type.i32)).output(0)
+    is_even = ov_opset.equal(x_mod, ov_opset.constant([0], Type.i32)).output(0)
+    med_index_0 = ov_opset.gather(x_sorted, ov_opset.floor(half_index).output(0), axis).output(0)   # COME BACK, does it sort out higher dimensions?
+    med_index_1 = ov_opset.gather(x_sorted, ov_opset.add(med_index_0, ov_opset.constant([1], Type.i32)).output(0), axis).output(0)
+    
+    median_odd = med_index_0
+    median_even = ov_opset.divide(ov_opset.add(med_index_1, med_index_0).output(0), ov_opset.constant([2], Type.i32))
+    
+    median_eval = ov_opset.select(is_even, median_even, median_odd)
+    
+    if keepdims == True:
+        if flattened == True:
+            median_shape = ov_opset.divide(x_shape_original, x_shape_original).output(0)
+            median_eval = ov_opset.reshape(median_eval, median_shape, False).output(0)
+        elif int_axis == True:
+            median_shape = ov_opset.shape_of(median_eval).output(0)
+            median_shape = ov_opset.unsqueeze(median_shape, axis).output(0)
+            median_eval = ov_opset.reshape(median_eval, median_shape, False).output(0)
+        else:
+            median_eval = ov_opset.unsqueeze(median_eval, ov_axis).output(0)
+        
+    return OpenVINOKerasTensor(median_eval)
 
 def meshgrid(*x, indexing="xy"):
     raise NotImplementedError(

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+env =
+    KERAS_BACKEND=openvino