keras-nightly 3.14.0.dev2026010104__py3-none-any.whl → 3.14.0.dev2026012204__py3-none-any.whl

keras/_tf_keras/keras/dtype_policies/__init__.py

@@ -7,6 +7,9 @@ since your modifications would be overwritten.
 from keras.src.dtype_policies import deserialize as deserialize
 from keras.src.dtype_policies import get as get
 from keras.src.dtype_policies import serialize as serialize
+from keras.src.dtype_policies.dtype_policy import (
+    AWQDTypePolicy as AWQDTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import DTypePolicy as DTypePolicy
 from keras.src.dtype_policies.dtype_policy import (
     FloatDTypePolicy as FloatDTypePolicy,

keras/_tf_keras/keras/ops/__init__.py

@@ -245,6 +245,7 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter
@@ -256,6 +257,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/_tf_keras/keras/ops/numpy/__init__.py

@@ -129,6 +129,7 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter
@@ -140,6 +141,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/_tf_keras/keras/quantizers/__init__.py

@@ -7,6 +7,7 @@ since your modifications would be overwritten.
 from keras.src.quantizers import deserialize as deserialize
 from keras.src.quantizers import get as get
 from keras.src.quantizers import serialize as serialize
+from keras.src.quantizers.awq_config import AWQConfig as AWQConfig
 from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantization_config import (
     Float8QuantizationConfig as Float8QuantizationConfig,

keras/dtype_policies/__init__.py

@@ -7,6 +7,9 @@ since your modifications would be overwritten.
 from keras.src.dtype_policies import deserialize as deserialize
 from keras.src.dtype_policies import get as get
 from keras.src.dtype_policies import serialize as serialize
+from keras.src.dtype_policies.dtype_policy import (
+    AWQDTypePolicy as AWQDTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import DTypePolicy as DTypePolicy
 from keras.src.dtype_policies.dtype_policy import (
     FloatDTypePolicy as FloatDTypePolicy,

keras/ops/__init__.py

@@ -245,6 +245,7 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter
@@ -256,6 +257,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/ops/numpy/__init__.py

@@ -129,6 +129,7 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
 from keras.src.ops.numpy import nextafter as nextafter
@@ -140,6 +141,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/quantizers/__init__.py

@@ -7,6 +7,7 @@ since your modifications would be overwritten.
 from keras.src.quantizers import deserialize as deserialize
 from keras.src.quantizers import get as get
 from keras.src.quantizers import serialize as serialize
+from keras.src.quantizers.awq_config import AWQConfig as AWQConfig
 from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantization_config import (
     Float8QuantizationConfig as Float8QuantizationConfig,

keras/src/backend/jax/nn.py

@@ -1471,25 +1471,42 @@ def _can_use_flash_attention(query, key, value, bias, raise_error=False):
         # Only support at least Ampere
         if not check_compute_capability("8.0"):
             raise RuntimeError("Require at least Ampere arch to run")
-        # Check inputs layout
+
+        # Inspect inputs of `check_layout`
         check_layout_params = list(
             inspect.signature(check_layout).parameters.keys()
         )
         for known_param in ("query", "key", "value", "bias", "layout"):
             check_layout_params.remove(known_param)
         # Defaults to `None` when not specified.
-        kwargs = {key: None for key in check_layout_params}
+        check_layout_kwargs = {key: None for key in check_layout_params}
         check_layout(
-            query, key, value, bias, layout=_normalize_layout("BTNH"), **kwargs
-        )
-        check_is_flash_attention(
             query,
             key,
-            _normalize_layout("BTNH"),
-            cudnn_version,
-            bias is not None,
-            is_training=False,
+            value,
+            bias,
+            layout=_normalize_layout("BTNH"),
+            **check_layout_kwargs,
         )
+
+        # Inspect inputs of `check_is_flash_attention`
+        check_is_flash_attention_params = inspect.signature(
+            check_is_flash_attention
+        ).parameters
+        check_is_flash_attention_kwargs = {
+            "query": query,
+            "key": key,
+            "value": value,
+            "layout": _normalize_layout("BTNH"),
+            "cudnn_version": cudnn_version,
+            "has_bias": bias is not None,
+            "is_training": False,
+        }
+        # Remove unsupported arguments
+        for param in list(check_is_flash_attention_kwargs.keys()):
+            if param not in check_is_flash_attention_params:
+                check_is_flash_attention_kwargs.pop(param)
+        check_is_flash_attention(**check_is_flash_attention_kwargs)
         return True
     except:
         if raise_error:

keras/src/backend/jax/numpy.py

@@ -1013,6 +1013,11 @@ def moveaxis(x, source, destination):
     return jnp.moveaxis(x, source=source, destination=destination)
 
 
+def nansum(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    return jnp.nansum(x, axis=axis, keepdims=keepdims)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     x = convert_to_tensor(x)
     return jnp.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)
@@ -1063,6 +1068,11 @@ def prod(x, axis=None, keepdims=False, dtype=None):
     return jnp.prod(x, axis=axis, keepdims=keepdims, dtype=dtype)
 
 
+def ptp(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    return jnp.ptp(x, axis=axis, keepdims=keepdims)
+
+
 def quantile(x, q, axis=None, method="linear", keepdims=False):
     x = convert_to_tensor(x)
     q = convert_to_tensor(q)

keras/src/backend/numpy/numpy.py

@@ -960,6 +960,17 @@ def moveaxis(x, source, destination):
     return np.moveaxis(x, source=source, destination=destination)
 
 
+def nansum(x, axis=None, keepdims=False):
+    axis = standardize_axis_for_numpy(axis)
+    dtype = standardize_dtype(x.dtype)
+
+    if dtype in ("bool", "int8", "int16"):
+        dtype = "int32"
+    elif dtype in ("uint8", "uint16"):
+        dtype = "uint32"
+    return np.nansum(x, axis=axis, keepdims=keepdims).astype(dtype)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     return np.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)
 
@@ -1018,6 +1029,10 @@ def prod(x, axis=None, keepdims=False, dtype=None):
     return np.prod(x, axis=axis, keepdims=keepdims, dtype=dtype)
 
 
+def ptp(x, axis=None, keepdims=False):
+    return np.ptp(x, axis=axis, keepdims=keepdims)
+
+
 def quantile(x, q, axis=None, method="linear", keepdims=False):
     axis = standardize_axis_for_numpy(axis)
     x = convert_to_tensor(x)

keras/src/backend/openvino/numpy.py

@@ -4,6 +4,7 @@ from openvino import Type
 
 from keras.src.backend import config
 from keras.src.backend.common import dtypes
+from keras.src.backend.common.backend_utils import canonicalize_axis
 from keras.src.backend.common.variables import standardize_dtype
 from keras.src.backend.openvino.core import DTYPES_MAX
 from keras.src.backend.openvino.core import DTYPES_MIN
@@ -705,7 +706,16 @@ def broadcast_to(x, shape):
 
 
 def cbrt(x):
-    raise NotImplementedError("`cbrt` is not supported with openvino backend")
+    x = get_ov_output(x)
+    x_type = x.get_element_type()
+    if x_type.is_integral() or x_type == Type.boolean:
+        x = ov_opset.convert(x, OPENVINO_DTYPES[config.floatx()]).output(0)
+    sign_x = ov_opset.sign(x)
+    abs_x = ov_opset.absolute(x)
+    one_third = ov_opset.constant(1.0 / 3.0, x.get_element_type())
+    root_abs = ov_opset.power(abs_x, one_third)
+    res = ov_opset.multiply(sign_x, root_abs)
+    return OpenVINOKerasTensor(res.output(0))
 
 
 def ceil(x):
@@ -893,9 +903,53 @@ def diag(x, k=0):
 
 
 def diagonal(x, offset=0, axis1=0, axis2=1):
-    raise NotImplementedError(
-        "`diagonal` is not supported with openvino backend"
-    )
+    x = get_ov_output(x)
+    shape = x.get_partial_shape()
+    rank = x.get_partial_shape().rank.get_length()
+    if rank is None:
+        raise ValueError("`diagonal` requires input tensor with static rank.")
+    if rank < 2:
+        raise ValueError(
+            f"diagonal requires input tensor with rank >= 2.Given rank: {rank}"
+        )
+    axis1 = canonicalize_axis(axis1, rank)
+    axis2 = canonicalize_axis(axis2, rank)
+    if axis1 == axis2:
+        raise ValueError("`axis1` and `axis2` cannot be the same.")
+
+    perm_order = [axis1, axis2] + [
+        i for i in range(rank) if i != axis1 and i != axis2
+    ]
+    perm_const = ov_opset.constant(perm_order, dtype=Type.i32).output(0)
+    x_transposed = ov_opset.transpose(x, perm_const)
+
+    N_dim = shape[axis1]
+    M_dim = shape[axis2]
+    if not N_dim.is_static or not M_dim.is_static:
+        raise ValueError(
+            "`diagonal` requires input tensor with static shape for axes "
+            f"`axis1` ({axis1}) and `axis2` ({axis2})."
+        )
+    N = N_dim.get_length()
+    M = M_dim.get_length()
+    if offset >= 0:
+        L = np.minimum(N, M - offset) if (M - offset) > 0 else 0
+        indices = [[i, i + offset] for i in range(L)]
+    else:
+        L = np.minimum(N + offset, M) if (N + offset) > 0 else 0
+        indices = [[i - offset, i] for i in range(L)]
+
+    indices = np.array(indices, dtype=np.int32).reshape(L, 2)
+    indices_const = ov_opset.constant(indices, dtype=Type.i32).output(0)
+
+    diag_gathered = ov_opset.gather_nd(x_transposed, indices_const)
+
+    out_rank = rank - 1
+    out_perm_order = list(range(1, out_rank)) + [0]
+    out_perm_const = ov_opset.constant(out_perm_order, dtype=Type.i32).output(0)
+
+    final_output = ov_opset.transpose(diag_gathered, out_perm_const)
+    return OpenVINOKerasTensor(final_output.output(0))
 
 
 def diff(a, n=1, axis=-1):
@@ -1072,22 +1126,28 @@ def expm1(x):
 
 def flip(x, axis=None):
     x_node = get_ov_output(x)
-    ndim = x.ndim
+
+    # Using OpenVINO tensor shape
+    ndim = len(x_node.get_partial_shape())
     if ndim is None:
         raise ValueError(
-            "The `flip` operation does not support tensors with dynamic rank"
+            "The `flip` operation does not support tensors with dynamic rank "
             "for the OpenVINO backend."
         )
+
     if axis is None:
         axis = list(range(ndim))
     elif isinstance(axis, int):
         axis = [axis]
+
     axis = [a + ndim if a < 0 else a for a in axis]
+
     begin = [0] * ndim
     end = [0] * ndim
     strides = [1] * ndim
     for a in axis:
         strides[a] = -1
+
     all_ones_mask = [1] * ndim
     result = ov_opset.strided_slice(
         data=x_node,
@@ -1100,6 +1160,61 @@ def flip(x, axis=None):
     return OpenVINOKerasTensor(result.output(0))
 
 
+def rot90(array, k=1, axes=(0, 1)):
+    """Rotate an array by 90 degrees in the plane specified by axes."""
+    array = get_ov_output(array)
+
+    if not isinstance(axes, (tuple, list)) or len(axes) != 2:
+        raise ValueError("axes must be a tuple of length 2")
+
+    shape = array.get_partial_shape()
+    ndim = shape.rank.get_length()
+    if ndim is None:
+        raise ValueError(
+            "`rot90` does not support tensors with dynamic rank "
+            "for the OpenVINO backend."
+        )
+
+    axis1 = canonicalize_axis(axes[0], ndim)
+    axis2 = canonicalize_axis(axes[1], ndim)
+
+    if axis1 == axis2:
+        raise ValueError("axes must be different")
+
+    k = k % 4
+    if k == 0:
+        return OpenVINOKerasTensor(array)
+
+    result = array
+
+    for _ in range(k):
+        # 1️ Transpose axis1 <-> axis2
+        perm = list(range(ndim))
+        perm[axis1], perm[axis2] = perm[axis2], perm[axis1]
+        perm_const = ov_opset.constant(perm, Type.i32).output(0)
+        result = ov_opset.transpose(result, perm_const).output(0)
+
+        # 2️ Reverse along axis1 using StridedSlice
+        begin = [0] * ndim
+        end = [0] * ndim
+        strides = [1] * ndim
+        strides[axis1] = -1
+
+        begin_mask = [1] * ndim
+        end_mask = [1] * ndim
+
+        result = ov_opset.strided_slice(
+            data=result,
+            begin=begin,
+            end=end,
+            strides=strides,
+            begin_mask=begin_mask,
+            end_mask=end_mask,
+        ).output(0)
+
+    return OpenVINOKerasTensor(result)
+
+
 def floor(x):
     x = get_ov_output(x)
     x_type = x.get_element_type()
@@ -1176,7 +1291,34 @@ def hstack(xs):
 
 
 def hypot(x1, x2):
-    raise NotImplementedError("`hypot` is not supported with openvino backend")
+    element_type = None
+    if isinstance(x1, OpenVINOKerasTensor):
+        element_type = x1.output.get_element_type()
+    if isinstance(x2, OpenVINOKerasTensor):
+        element_type = x2.output.get_element_type()
+    x1 = get_ov_output(x1, element_type)
+    x2 = get_ov_output(x2, element_type)
+    x1, x2 = _align_operand_types(x1, x2, "hypot()")
+    x_type = x1.get_element_type()
+    if x_type.is_integral() or x_type == Type.boolean:
+        ov_type = OPENVINO_DTYPES[config.floatx()]
+        x1 = ov_opset.convert(x1, ov_type)
+        x2 = ov_opset.convert(x2, ov_type)
+    x1_abs = ov_opset.absolute(x1)
+    x2_abs = ov_opset.absolute(x2)
+    max_val = ov_opset.maximum(x1_abs, x2_abs)
+    min_val = ov_opset.minimum(x1_abs, x2_abs)
+    one = ov_opset.constant(1, max_val.get_element_type())
+    is_zero_mask = ov_opset.equal(
+        max_val, ov_opset.constant(0, max_val.get_element_type())
+    )
+    safe_divisor = ov_opset.select(is_zero_mask, one, max_val)
+    ratio = ov_opset.divide(min_val, safe_divisor)
+    result = ov_opset.multiply(
+        max_val,
+        ov_opset.sqrt(ov_opset.add(one, ov_opset.multiply(ratio, ratio))),
+    )
+    return OpenVINOKerasTensor(result.output(0))
 
 
 def identity(n, dtype=None):
@@ -1313,7 +1455,66 @@ def isreal(x):
 
 
 def kron(x1, x2):
-    raise NotImplementedError("`kron` is not supported with openvino backend")
+    x1 = get_ov_output(x1)
+    x2 = get_ov_output(x2)
+    x1, x2 = _align_operand_types(x1, x2, "kron()")
+    x1_shape = x1.get_partial_shape()
+    x2_shape = x2.get_partial_shape()
+    if x1_shape.rank.is_dynamic or x2_shape.rank.is_dynamic:
+        raise ValueError(
+            "`kron` does not support tensors with dynamic rank for "
+            "the OpenVINO backend."
+        )
+    ndim1 = x1_shape.rank.get_length()
+    ndim2 = x2_shape.rank.get_length()
+    if ndim1 < ndim2:
+        axes = ov_opset.range(
+            ov_opset.constant(0, Type.i32),
+            ov_opset.constant(ndim2 - ndim1, Type.i32),
+            ov_opset.constant(1, Type.i32),
+        )
+        x1 = ov_opset.unsqueeze(x1, axes)
+        ndim1 = ndim2
+    elif ndim2 < ndim1:
+        axes = ov_opset.range(
+            ov_opset.constant(0, Type.i32),
+            ov_opset.constant(ndim1 - ndim2, Type.i32),
+            ov_opset.constant(1, Type.i32),
+        )
+        x2 = ov_opset.unsqueeze(x2, axes)
+        ndim2 = ndim1
+    shape1 = ov_opset.shape_of(x1, Type.i32)
+    shape2 = ov_opset.shape_of(x2, Type.i32)
+    ones = ov_opset.broadcast(
+        ov_opset.constant(1, Type.i32), ov_opset.constant([ndim1], Type.i32)
+    )
+    axis = ov_opset.constant(1, Type.i32)
+    flatten = ov_opset.constant([-1], Type.i32)
+    unsqueezed_ones = ov_opset.unsqueeze(ones, axis)
+    x1_new_shape = ov_opset.reshape(
+        ov_opset.concat(
+            [ov_opset.unsqueeze(shape1, axis), unsqueezed_ones],
+            axis=1,
+        ),
+        flatten,
+        False,
+    )
+    x2_new_shape = ov_opset.reshape(
+        ov_opset.concat(
+            [unsqueezed_ones, ov_opset.unsqueeze(shape2, axis)],
+            axis=1,
+        ),
+        flatten,
+        False,
+    )
+    result = ov_opset.multiply(
+        ov_opset.reshape(x1, x1_new_shape, False),
+        ov_opset.reshape(x2, x2_new_shape, False),
+    )
+    result = ov_opset.reshape(
+        result, ov_opset.multiply(shape1, shape2), False
+    ).output(0)
+    return OpenVINOKerasTensor(result)
 
 
 def lcm(x1, x2):
@@ -1578,9 +1779,42 @@ def logaddexp(x1, x2):
 
 
 def logaddexp2(x1, x2):
-    raise NotImplementedError(
-        "`logaddexp2` is not supported with openvino backend"
+    element_type = None
+    if isinstance(x1, OpenVINOKerasTensor):
+        element_type = x1.output.get_element_type()
+    if isinstance(x2, OpenVINOKerasTensor):
+        element_type = x2.output.get_element_type()
+    x1 = get_ov_output(x1, element_type)
+    x2 = get_ov_output(x2, element_type)
+    x1, x2 = _align_operand_types(x1, x2, "logaddexp2()")
+
+    if x1.element_type.is_integral() or x2.element_type.is_integral():
+        float_dtype = OPENVINO_DTYPES[config.floatx()]
+        if x1.get_element_type().is_integral():
+            x1 = ov_opset.convert(x1, float_dtype)
+        if x2.get_element_type().is_integral():
+            x2 = ov_opset.convert(x2, float_dtype)
+
+    max_val = ov_opset.maximum(x1, x2)
+
+    sub = ov_opset.subtract(x1, x2)
+    abs_diff = ov_opset.abs(sub)
+
+    neg_abs_diff = ov_opset.negative(abs_diff)
+
+    element_type = neg_abs_diff.get_element_type()
+
+    two = ov_opset.constant(2, dtype=element_type)
+
+    power_of_2 = ov_opset.power(two, neg_abs_diff)
+
+    one_plus_power = ov_opset.add(
+        ov_opset.constant(1, dtype=element_type), power_of_2
     )
+    log2_term = ov_opset.divide(ov_opset.log(one_plus_power), ov_opset.log(two))
+    result = ov_opset.add(max_val, log2_term).output(0)
+
+    return OpenVINOKerasTensor(result)
 
 
 def logical_and(x1, x2):
@@ -1855,6 +2089,10 @@ def moveaxis(x, source, destination):
     return OpenVINOKerasTensor(ov_opset.transpose(x, axes_const).output(0))
 
 
+def nansum(x, axis=None, keepdims=False):
+    raise NotImplementedError("`nansum` is not supported with openvino backend")
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     x = get_ov_output(x)
     dtype = x.get_element_type()
@@ -2005,6 +2243,10 @@ def prod(x, axis=None, keepdims=False, dtype=None):
     return OpenVINOKerasTensor(result)
 
 
+def ptp(x, axis=None, keepdims=False):
+    raise NotImplementedError("`ptp` is not supported with openvino backend")
+
+
 def quantile(x, q, axis=None, method="linear", keepdims=False):
     raise NotImplementedError(
         "`quantile` is not supported with openvino backend"
@@ -2141,7 +2383,14 @@ def sinh(x):
 
 
 def size(x):
-    raise NotImplementedError("`size` is not supported with openvino backend")
+    x = get_ov_output(x)
+    shape_tensor = ov_opset.shape_of(x, output_type=Type.i64)
+    final_size = ov_opset.reduce_prod(
+        shape_tensor,
+        ov_opset.constant([0], Type.i64),
+        keep_dims=False,
+    )
+    return OpenVINOKerasTensor(final_size.output(0))
 
 
 def sort(x, axis=-1):
@@ -2283,9 +2532,20 @@ def std(x, axis=None, keepdims=False):
 
 
 def swapaxes(x, axis1, axis2):
-    raise NotImplementedError(
-        "`swapaxes` is not supported with openvino backend"
-    )
+    x = get_ov_output(x)
+    x_shape = x.get_partial_shape()
+    if x_shape.rank.is_dynamic:
+        raise ValueError(
+            "`swapaxes` does not support tensors with dynamic rank for the "
+            "OpenVINO backend."
+        )
+    rank = x_shape.rank.get_length()
+    axis1 = canonicalize_axis(axis1, rank)
+    axis2 = canonicalize_axis(axis2, rank)
+    axes = list(range(rank))
+    axes[axis1], axes[axis2] = axes[axis2], axes[axis1]
+    result = ov_opset.transpose(x, ov_opset.constant(axes, Type.i32))
+    return OpenVINOKerasTensor(result.output(0))
 
 
 def take(x, indices, axis=None):
@@ -2404,7 +2664,8 @@ def tile(x, repeats):
 
 
 def trace(x, offset=0, axis1=0, axis2=1):
-    raise NotImplementedError("`trace` is not supported with openvino backend")
+    x = diagonal(x, offset=offset, axis1=axis1, axis2=axis2)
+    return sum(x, axis=-1)
 
 
 def tri(N, M=None, k=0, dtype=None):
@@ -2937,6 +3198,66 @@ def slogdet(x):
 
 
 def argpartition(x, kth, axis=-1):
-    raise NotImplementedError(
-        "`argpartition` is not supported with openvino backend"
+    x = get_ov_output(x)
+    x_shape = x.get_partial_shape()
+    rank = x_shape.rank.get_length()
+    axis = canonicalize_axis(axis, rank)
+    axes = list(range(rank))
+    axes[axis], axes[-1] = axes[-1], axes[axis]
+    x = ov_opset.transpose(x, ov_opset.constant(axes))
+    x_shape_tensor = ov_opset.shape_of(x)
+    n = ov_opset.gather(
+        x_shape_tensor,
+        ov_opset.constant(-1),
+        ov_opset.constant(0),
+    )
+    if isinstance(kth, int) and kth < 0:
+        kth_tensor = ov_opset.add(
+            n,
+            ov_opset.constant(kth, n.get_element_type()),
+        )
+    else:
+        kth_tensor = ov_opset.constant(kth, n.get_element_type())
+    one = ov_opset.constant(1, kth_tensor.get_element_type())
+    k_val = ov_opset.add(kth_tensor, one)
+    bottom_ind = ov_opset.topk(
+        ov_opset.negative(x),
+        k=k_val,
+        axis=-1,
+        mode="max",
+        sort="value",
+    ).output(1)
+    one_hot_mask = ov_opset.one_hot(
+        bottom_ind,
+        n,
+        ov_opset.constant(1),
+        ov_opset.constant(0),
+        axis=-1,
     )
+    mask = ov_opset.reduce_sum(
+        one_hot_mask,
+        ov_opset.constant([-2]),
+        keep_dims=False,
+    )
+    ones = ov_opset.broadcast(
+        ov_opset.constant(1),
+        x_shape_tensor,
+    )
+    proxy = ov_opset.subtract(ones, mask)
+    remaining_k = ov_opset.subtract(n, k_val)
+    top_ind = ov_opset.topk(
+        proxy,
+        k=remaining_k,
+        axis=-1,
+        mode="max",
+        sort="value",
+    ).output(1)
+    result = ov_opset.concat([bottom_ind, top_ind], axis=-1)
+    inv_axes = [0] * rank
+    for i, a in enumerate(axes):
+        inv_axes[a] = i
+    result = ov_opset.transpose(
+        result,
+        ov_opset.constant(inv_axes),
+    ).output(0)
+    return OpenVINOKerasTensor(result)