fix

Author: Saurabh Singh

keras/src/backend/openvino/numpy.py

@@ -996,64 +996,50 @@ def maximum(x1, x2):
 
 
 def median(x, axis=None, keepdims=False):
-    x = get_ov_output(x)
-    
+    x_node = get_ov_output(x)
+    orig_dtype = x_node.get_element_type()
+
     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
-    
-    shape = ov_opset.shape_of(x).output(0)
-    shape = ov_opset.convert(shape, Type.i64).output(0)
-    
-    if axis is not None:
-        indices = ov_opset.constant([axis], Type.i32).output(0)
-        length = ov_opset.gather(shape, indices, 0).output(0)
-        length = ov_opset.reshape(length, ov_opset.constant([], Type.i32).output(0), False).output(0)
+        flatten_shape = ov_opset.constant([-1], Type.i64).output(0)
+        x_node = ov_opset.reshape(x_node, flatten_shape, False).output(0)
+        axis_val = 0
     else:
-        length = ov_opset.shape_of(shape).output(0)
-        length = ov_opset.convert(length, Type.i64).output(0)
-        length = ov_opset.reshape(length, ov_opset.constant([], Type.i32).output(0), False).output(0)
-    
-    sorted_x = ov_opset.topk(x, length, axis, "min", "value").output(0)
-    
-    const_2 = ov_opset.constant(2, Type.i64).output(0)
-    mid_index = ov_opset.floor_mod(length, const_2).output(0)
-    is_odd = ov_opset.equal(mid_index, ov_opset.constant(1, Type.i64).output(0)).output(0)
-    
-    half_length = ov_opset.divide(length, const_2).output(0)
-    floor_half_length = ov_opset.floor(half_length).output(0)
-    floor_half_length = ov_opset.convert(floor_half_length, Type.i64).output(0)
-    
-    mid_index_scalar = ov_opset.convert(floor_half_length, Type.i32).output(0)
-    middle_elem = ov_opset.gather(sorted_x, mid_index_scalar, axis).output(0)
-    
-    prev_mid_index = ov_opset.subtract(floor_half_length, ov_opset.constant(1, Type.i64).output(0)).output(0)
-    prev_mid_index_scalar = ov_opset.convert(prev_mid_index, Type.i32).output(0)
-    prev_middle_elem = ov_opset.gather(sorted_x, prev_mid_index_scalar, axis).output(0)
-    
-    median_value = ov_opset.select(
-        is_odd,
-        middle_elem,
-        ov_opset.divide(
-            ov_opset.add(middle_elem, prev_middle_elem).output(0),
-            ov_opset.constant(2.0, middle_elem.get_element_type()).output(0)
-        ).output(0)
-    ).output(0)
-    
+        axis_val = axis
+
+    shape = ov_opset.shape_of(x_node).output(0)
+    shape = ov_opset.convert(shape, Type.i64).output(0)
+
+    axis_const = ov_opset.constant([axis_val], Type.i64).output(0)
+    length = ov_opset.gather(shape, axis_const, 0).output(0)
+    length = ov_opset.reshape(length, ov_opset.constant([], Type.i64).output(0), False).output(0)
+
+    sorted_x = ov_opset.topk(x_node, length, axis_val, "min", "value").output(0)
+
+    two = ov_opset.constant(2, Type.i64).output(0)
+
+    half = ov_opset.divide(length, two).output(0)
+    mid_idx = ov_opset.floor(half).output(0)
+    rem = ov_opset.floor_mod(length, two).output(0)
+    is_odd = ov_opset.equal(rem, ov_opset.constant(1, Type.i64).output(0)).output(0)
+
+    mid_idx_i32 = ov_opset.convert(mid_idx, Type.i32).output(0)
+    middle = ov_opset.gather(sorted_x, mid_idx_i32, axis_val).output(0)
+
+    prev_idx = ov_opset.subtract(mid_idx, ov_opset.constant(1, Type.i64).output(0)).output(0)
+    prev_idx_i32 = ov_opset.convert(prev_idx, Type.i32).output(0)
+    prev = ov_opset.gather(sorted_x, prev_idx_i32, axis_val).output(0)
+    sum_val = ov_opset.add(middle, prev).output(0)
+    avg = ov_opset.divide(sum_val, ov_opset.constant(2.0, sum_val.get_element_type()).output(0)).output(0)
+
+    median_val = ov_opset.select(is_odd, middle, avg).output(0)
+
     if keepdims:
-        keep_shape = shape
-        if axis is not None:
-            one_tensor = ov_opset.constant(1, Type.i64).output(0)
-            indices = ov_opset.constant([axis], Type.i32).output(0)
-            keep_shape = ov_opset.scatter_elements_update(shape, indices, one_tensor, 0).output(0)
-        median_value = ov_opset.reshape(median_value, keep_shape, False).output(0)
-    elif axis is None and x.get_partial_shape().rank.get_length() > 1:
-        scalar_shape = ov_opset.constant([], Type.i32).output(0)
-        median_value = ov_opset.reshape(median_value, scalar_shape, False).output(0)
+        one = ov_opset.constant(1, Type.i64).output(0)
+        keep_shape = ov_opset.scatter_elements_update(shape, axis_const, one, 0).output(0)
+        median_val = ov_opset.reshape(median_val, keep_shape, False).output(0)
 
-    median_value = ov_opset.convert(median_value, x.get_element_type()).output(0)
-    return OpenVINOKerasTensor(median_value)
+    median_val = ov_opset.convert(median_val, orig_dtype).output(0)
+    return OpenVINOKerasTensor(median_val)
 
 
 def meshgrid(*x, indexing="xy"):