keras-nightly 3.12.0.dev2025100503__py3-none-any.whl → 3.14.0.dev2026011604__py3-none-any.whl

keras/src/backend/jax/numpy.py

@@ -446,6 +446,11 @@ def array(x, dtype=None):
     return jnp.array(x, dtype=dtype)
 
 
+def view(x, dtype=None):
+    x = convert_to_tensor(x)
+    return x.view(dtype=dtype)
+
+
 def average(x, axis=None, weights=None):
     x = convert_to_tensor(x)
     dtypes_to_resolve = [x.dtype, float]
@@ -673,6 +678,10 @@ def empty(shape, dtype=None):
     return jnp.empty(shape, dtype=dtype)
 
 
+def empty_like(x, dtype=None):
+    return jnp.empty_like(x, dtype=dtype)
+
+
 def equal(x1, x2):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
@@ -819,6 +828,11 @@ def isposinf(x):
     return jnp.isposinf(x)
 
 
+def isreal(x):
+    x = convert_to_tensor(x)
+    return jnp.isreal(x)
+
+
 def kron(x1, x2):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
@@ -831,6 +845,19 @@ def lcm(x1, x2):
     return jnp.lcm(x1, x2)
 
 
+def ldexp(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    if standardize_dtype(x2.dtype) not in dtypes.INT_TYPES:
+        raise TypeError(
+            f"ldexp exponent must be an integer type. "
+            f"Received: x2 dtype={x2.dtype}"
+        )
+
+    return jnp.ldexp(x1, x2)
+
+
 def less(x1, x2):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
@@ -1036,6 +1063,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)
@@ -1157,6 +1189,11 @@ def split(x, indices_or_sections, axis=0):
     return jnp.split(x, indices_or_sections, axis=axis)
 
 
+def array_split(x, indices_or_sections, axis=0):
+    x = convert_to_tensor(x)
+    return jnp.array_split(x, indices_or_sections, axis=axis)
+
+
 def stack(x, axis=0):
     x = [convert_to_tensor(t) for t in x]
     return jnp.stack(x, axis=axis)
@@ -1181,6 +1218,8 @@ def take(x, indices, axis=None):
 
 
 def take_along_axis(x, indices, axis=None):
+    x = convert_to_tensor(x)
+    indices = convert_to_tensor(indices, sparse=False)
     return jnp.take_along_axis(x, indices, axis=axis)
 
 
@@ -1235,14 +1274,7 @@ def tile(x, repeats):
 
 def trace(x, offset=0, axis1=0, axis2=1):
     x = convert_to_tensor(x)
-    dtype = None
-    # TODO: Remove the condition of uint8 and uint16 once we have jax>=0.4.27
-    # for both CPU & GPU environments.
-    # uint8 and uint16 will be casted to uint32 when jax>=0.4.27 but to int32
-    # otherwise.
-    if standardize_dtype(x.dtype) in ("bool", "uint8", "uint16"):
-        dtype = "int32"
-    return jnp.trace(x, offset=offset, axis1=axis1, axis2=axis2, dtype=dtype)
+    return jnp.trace(x, offset=offset, axis1=axis1, axis2=axis2)
 
 
 def tri(N, M=None, k=0, dtype=None):
@@ -1324,6 +1356,12 @@ def negative(x):
     return jnp.negative(x)
 
 
+def nextafter(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    return jnp.nextafter(x1, x2)
+
+
 @sparse.elementwise_unary(linear=False)
 def square(x):
     x = convert_to_tensor(x)
@@ -1363,6 +1401,19 @@ def transpose(x, axes=None):
     return jnp.transpose(x, axes=axes)
 
 
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    y = convert_to_tensor(y)
+    if x is not None:
+        x = convert_to_tensor(x)
+    dx = convert_to_tensor(dx)
+    return jnp.trapezoid(y, x, dx=dx, axis=axis)
+
+
+def vander(x, N=None, increasing=False):
+    x = convert_to_tensor(x)
+    return jnp.vander(x, N=N, increasing=increasing)
+
+
 def var(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     # `jnp.var` does not handle low precision (e.g., float16) overflow

keras/src/backend/jax/trainer.py

@@ -266,8 +266,14 @@ class JAXTrainer(base_trainer.Trainer):
             if distribution_lib.distribution() is not None:
                 state_shardings = self._get_state_sharding_spec()
                 out_shardings = (None, state_shardings)
+            if is_nnx_enabled():
+                step_fn = lambda state, data: type(self).train_step(
+                    self, state, data
+                )
+            else:
+                step_fn = self.train_step
             train_step = jit(
-                self.train_step,
+                step_fn,
                 donate_argnums=0,
                 out_shardings=out_shardings,
             )
@@ -296,8 +302,14 @@ class JAXTrainer(base_trainer.Trainer):
                     metrics_shardings,
                 )
                 out_shardings = (None, state_shardings)
+            if is_nnx_enabled():
+                step_fn = lambda state, data: type(self).test_step(
+                    self, state, data
+                )
+            else:
+                step_fn = self.test_step
             test_step = jit(
-                self.test_step,
+                step_fn,
                 donate_argnums=0,
                 out_shardings=out_shardings,
             )

keras/src/backend/numpy/linalg.py

@@ -96,3 +96,7 @@ def lstsq(a, b, rcond=None):
     a = convert_to_tensor(a)
     b = convert_to_tensor(b)
     return np.linalg.lstsq(a, b, rcond=rcond)[0]
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    raise NotImplementedError("JVP is not supported by the Numpy backend.")

keras/src/backend/numpy/nn.py

@@ -3,6 +3,9 @@ import numpy as np
 from jax import lax
 
 from keras.src import backend
+from keras.src.backend.common.backend_utils import (
+    compute_adaptive_pooling_window_sizes,
+)
 from keras.src.backend.common.backend_utils import (
     compute_conv_transpose_padding_args_for_jax,
 )
@@ -340,6 +343,252 @@ def average_pool(
         return pooled / window_counts
 
 
+def _compute_adaptive_pooling_gather_indices(
+    input_dim, output_size, big_window
+):
+    window_starts = np.floor(
+        (np.arange(output_size) * input_dim) / output_size
+    ).astype(np.int32)
+
+    window_ends = np.ceil(
+        (np.arange(1, output_size + 1) * input_dim) / output_size
+    ).astype(np.int32)
+
+    window_sizes = window_ends - window_starts
+    is_big = window_sizes == big_window
+
+    small_window = big_window - 1
+    small_pool_len = input_dim - small_window + 1
+
+    small_indices = window_starts
+    big_indices = window_starts + small_pool_len
+
+    gather = np.where(is_big, big_indices, small_indices)
+    return gather.astype(np.int32)
+
+
+def _strided_view_1d(x, window_size):
+    n, l, c = x.shape
+    out = l - window_size + 1
+
+    strides = x.strides
+    shape = (n, out, window_size, c)
+    new_strides = (strides[0], strides[1], strides[1], strides[2])
+
+    return np.lib.stride_tricks.as_strided(x, shape=shape, strides=new_strides)
+
+
+def _adaptive_pool1d_impl(inputs, output_size, mode, data_format):
+    if isinstance(output_size, int):
+        output_size = (output_size,)
+
+    if data_format == "channels_first":
+        inputs = np.transpose(inputs, (0, 2, 1))
+
+    n, l, c = inputs.shape
+    out_l = output_size[0]
+
+    small, big = compute_adaptive_pooling_window_sizes(l, out_l)
+    gather = _compute_adaptive_pooling_gather_indices(l, out_l, big)
+
+    sv_small = _strided_view_1d(inputs, small)
+    small_pool = (
+        np.mean(sv_small, axis=2)
+        if mode == "average"
+        else np.max(sv_small, axis=2)
+    )
+
+    sv_big = _strided_view_1d(inputs, big)
+    big_pool = (
+        np.mean(sv_big, axis=2) if mode == "average" else np.max(sv_big, axis=2)
+    )
+
+    combined = np.concatenate([small_pool, big_pool], axis=1)
+    out = combined[:, gather, :]
+
+    if data_format == "channels_first":
+        out = np.transpose(out, (0, 2, 1))
+
+    return out
+
+
+def _adaptive_pool2d_impl(inputs, output_size, mode, data_format):
+    if isinstance(output_size, int):
+        output_size = (output_size, output_size)
+
+    if data_format == "channels_first":
+        inputs = np.transpose(inputs, (0, 2, 3, 1))
+
+    n, h, w, c = inputs.shape
+    out_h, out_w = output_size
+
+    small_h, big_h = compute_adaptive_pooling_window_sizes(h, out_h)
+    gather_h = _compute_adaptive_pooling_gather_indices(h, out_h, big_h)
+
+    x_h = np.transpose(inputs, (0, 2, 1, 3)).reshape(n * w, h, c)
+
+    sv_small_h = _strided_view_1d(x_h, small_h)
+    small_pool_h = (
+        np.mean(sv_small_h, axis=2)
+        if mode == "average"
+        else np.max(sv_small_h, axis=2)
+    )
+
+    sv_big_h = _strided_view_1d(x_h, big_h)
+    big_pool_h = (
+        np.mean(sv_big_h, axis=2)
+        if mode == "average"
+        else np.max(sv_big_h, axis=2)
+    )
+
+    combined_h = np.concatenate([small_pool_h, big_pool_h], axis=1)
+    pooled_h = combined_h[:, gather_h, :]
+
+    pooled_h = pooled_h.reshape(n, w, out_h, c)
+    pooled_h = np.transpose(pooled_h, (0, 2, 1, 3))
+
+    small_w, big_w = compute_adaptive_pooling_window_sizes(w, out_w)
+    gather_w = _compute_adaptive_pooling_gather_indices(w, out_w, big_w)
+
+    x_w = pooled_h.reshape(n * out_h, w, c)
+
+    sv_small_w = _strided_view_1d(x_w, small_w)
+    small_pool_w = (
+        np.mean(sv_small_w, axis=2)
+        if mode == "average"
+        else np.max(sv_small_w, axis=2)
+    )
+
+    sv_big_w = _strided_view_1d(x_w, big_w)
+    big_pool_w = (
+        np.mean(sv_big_w, axis=2)
+        if mode == "average"
+        else np.max(sv_big_w, axis=2)
+    )
+
+    combined_w = np.concatenate([small_pool_w, big_pool_w], axis=1)
+    out = combined_w[:, gather_w, :].reshape(n, out_h, out_w, c)
+
+    if data_format == "channels_first":
+        out = np.transpose(out, (0, 3, 1, 2))
+
+    return out
+
+
+def _adaptive_pool3d_impl(inputs, output_size, mode, data_format):
+    if isinstance(output_size, int):
+        output_size = (output_size, output_size, output_size)
+
+    if data_format == "channels_first":
+        inputs = np.transpose(inputs, (0, 2, 3, 4, 1))
+
+    n, d, h, w, c = inputs.shape
+    out_d, out_h, out_w = output_size
+
+    small_d, big_d = compute_adaptive_pooling_window_sizes(d, out_d)
+    gather_d = _compute_adaptive_pooling_gather_indices(d, out_d, big_d)
+
+    x_d = np.transpose(inputs, (0, 2, 3, 1, 4)).reshape(n * h * w, d, c)
+
+    sv_small_d = _strided_view_1d(x_d, small_d)
+    small_pool_d = (
+        np.mean(sv_small_d, axis=2)
+        if mode == "average"
+        else np.max(sv_small_d, axis=2)
+    )
+
+    sv_big_d = _strided_view_1d(x_d, big_d)
+    big_pool_d = (
+        np.mean(sv_big_d, axis=2)
+        if mode == "average"
+        else np.max(sv_big_d, axis=2)
+    )
+
+    combined_d = np.concatenate([small_pool_d, big_pool_d], axis=1)
+    pooled_d = combined_d[:, gather_d, :].reshape(n, h, w, out_d, c)
+    pooled_d = np.transpose(pooled_d, (0, 3, 1, 2, 4))
+
+    small_h, big_h = compute_adaptive_pooling_window_sizes(h, out_h)
+    gather_h = _compute_adaptive_pooling_gather_indices(h, out_h, big_h)
+
+    x_h = np.transpose(pooled_d, (0, 1, 3, 2, 4)).reshape(n * out_d * w, h, c)
+
+    sv_small_h = _strided_view_1d(x_h, small_h)
+    small_pool_h = (
+        np.mean(sv_small_h, axis=2)
+        if mode == "average"
+        else np.max(sv_small_h, axis=2)
+    )
+
+    sv_big_h = _strided_view_1d(x_h, big_h)
+    big_pool_h = (
+        np.mean(sv_big_h, axis=2)
+        if mode == "average"
+        else np.max(sv_big_h, axis=2)
+    )
+
+    combined_h = np.concatenate([small_pool_h, big_pool_h], axis=1)
+    pooled_h = combined_h[:, gather_h, :].reshape(n, out_d, w, out_h, c)
+    pooled_h = np.transpose(pooled_h, (0, 1, 3, 2, 4))
+
+    small_w, big_w = compute_adaptive_pooling_window_sizes(w, out_w)
+    gather_w = _compute_adaptive_pooling_gather_indices(w, out_w, big_w)
+
+    x_w = pooled_h.reshape(n * out_d * out_h, w, c)
+
+    sv_small_w = _strided_view_1d(x_w, small_w)
+    small_pool_w = (
+        np.mean(sv_small_w, axis=2)
+        if mode == "average"
+        else np.max(sv_small_w, axis=2)
+    )
+
+    sv_big_w = _strided_view_1d(x_w, big_w)
+    big_pool_w = (
+        np.mean(sv_big_w, axis=2)
+        if mode == "average"
+        else np.max(sv_big_w, axis=2)
+    )
+
+    combined_w = np.concatenate([small_pool_w, big_pool_w], axis=1)
+    out = combined_w[:, gather_w, :].reshape(n, out_d, out_h, out_w, c)
+
+    if data_format == "channels_first":
+        out = np.transpose(out, (0, 4, 1, 2, 3))
+
+    return out
+
+
+def adaptive_average_pool(inputs, output_size, data_format=None):
+    data_format = backend.standardize_data_format(data_format)
+    dims = inputs.ndim - 2
+    if dims == 1:
+        return _adaptive_pool1d_impl(
+            inputs, output_size, "average", data_format
+        )
+    if dims == 2:
+        return _adaptive_pool2d_impl(
+            inputs, output_size, "average", data_format
+        )
+    if dims == 3:
+        return _adaptive_pool3d_impl(
+            inputs, output_size, "average", data_format
+        )
+    raise ValueError("adaptive_average_pool supports only 1D/2D/3D")
+
+
+def adaptive_max_pool(inputs, output_size, data_format=None):
+    data_format = backend.standardize_data_format(data_format)
+    dims = inputs.ndim - 2
+    if dims == 1:
+        return _adaptive_pool1d_impl(inputs, output_size, "max", data_format)
+    if dims == 2:
+        return _adaptive_pool2d_impl(inputs, output_size, "max", data_format)
+    if dims == 3:
+        return _adaptive_pool3d_impl(inputs, output_size, "max", data_format)
+    raise ValueError("adaptive_max_pool supports only 1D/2D/3D")
+
+
 def _convert_to_lax_conv_dimension_numbers(
     num_spatial_dims,
     data_format="channels_last",
@@ -404,7 +653,7 @@ def conv(
             f"kernel in_channels {kernel_in_channels}. "
         )
     feature_group_count = channels // kernel_in_channels
-    return np.array(
+    result = np.array(
         jax.lax.conv_general_dilated(
             inputs,
             kernel if is_tensor(kernel) else kernel.numpy(),
@@ -415,6 +664,14 @@ def conv(
             feature_group_count=feature_group_count,
         )
     )
+    if result.size == 0:
+        raise ValueError(
+            "The convolution operation resulted in an empty output. "
+            "This can happen if the input is too small for the given "
+            "kernel size, strides, dilation rate, and padding mode. "
+            "Please check the input shape and convolution parameters."
+        )
+    return result
 
 
 def depthwise_conv(
@@ -1176,3 +1433,56 @@ def dot_product_attention(
     return _dot_product_attention_xla(
         query, key, value, bias, mask, is_causal, scale
     )
+
+
+def unfold(input, kernel_size, dilation=1, padding=0, stride=1):
+    """NumPy implementation of Unfold.
+    Extract sliding local blocks from a **NCHW** batched image tensor.
+
+    Args:
+        input: 4-D tensor, shape (N, C, H, W)  **required**.
+        kernel_size: int or (kH, kW)
+        dilation: int or (dH, dW), default 1
+        padding: int or (pH, pW), default 0
+        stride: int or (sH, sW), default 1
+
+    Returns:
+        3-D tensor, shape (N, C*kH*kW, L)
+    """
+
+    def _pair(x):
+        return (x, x) if isinstance(x, int) else x
+
+    k = _pair(kernel_size)
+    d = _pair(dilation)
+    p = _pair(padding)
+    s = _pair(stride)
+
+    N, C, H, W = input.shape
+
+    # ---- padding ----
+    if any(_ > 0 for _ in p):
+        input = np.pad(
+            input, ((0, 0), (0, 0), (p[0], p[0]), (p[1], p[1])), mode="constant"
+        )
+
+    # ----  spatial size ----
+    oH = (input.shape[2] - (k[0] - 1) * d[0] - 1) // s[0] + 1
+    oW = (input.shape[3] - (k[1] - 1) * d[1] - 1) // s[1] + 1
+
+    i0 = np.arange(0, oH) * s[0]
+    j0 = np.arange(0, oW) * s[1]
+    i, j = np.meshgrid(i0, j0, indexing="ij")  # shape (oH, oW)
+    i = i.reshape(-1)
+    j = j.reshape(-1)
+
+    # ---- flatten patches ----
+    patches = np.empty((N, C, k[0], k[1], oH * oW), dtype=input.dtype)
+    for idx in range(k[0]):
+        for jdx in range(k[1]):
+            patches[:, :, idx, jdx, :] = input[
+                :, :, i + idx * d[0], j + jdx * d[1]
+            ]
+
+    # ---- reshape -> (N, C*kH*kW, L) ----
+    return patches.reshape(N, C * k[0] * k[1], -1)

keras/src/backend/numpy/numpy.py

@@ -294,6 +294,11 @@ def array(x, dtype=None):
     return convert_to_tensor(x, dtype=dtype)
 
 
+def view(x, dtype=None):
+    x = convert_to_tensor(x)
+    return x.view(dtype=dtype)
+
+
 def average(x, axis=None, weights=None):
     axis = standardize_axis_for_numpy(axis)
     x = convert_to_tensor(x)
@@ -607,6 +612,10 @@ def empty(shape, dtype=None):
     return np.empty(shape, dtype=dtype)
 
 
+def empty_like(x, dtype=None):
+    return np.empty_like(x, dtype=dtype)
+
+
 def equal(x1, x2):
     return np.equal(x1, x2)
 
@@ -745,6 +754,11 @@ def isposinf(x):
     return np.isposinf(x)
 
 
+def isreal(x):
+    x = convert_to_tensor(x)
+    return np.isreal(x)
+
+
 def kron(x1, x2):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
@@ -759,6 +773,19 @@ def lcm(x1, x2):
     return np.lcm(x1, x2).astype(dtype)
 
 
+def ldexp(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
+
+    if standardize_dtype(x2.dtype) not in dtypes.INT_TYPES:
+        raise TypeError(
+            f"ldexp exponent must be an integer type. "
+            f"Received: x2 dtype={x2.dtype}"
+        )
+    return np.ldexp(x1, x2).astype(dtype)
+
+
 def less(x1, x2):
     return np.less(x1, x2)
 
@@ -991,6 +1018,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)
@@ -1097,6 +1128,11 @@ def split(x, indices_or_sections, axis=0):
     return np.split(x, indices_or_sections, axis=axis)
 
 
+def array_split(x, indices_or_sections, axis=0):
+    axis = standardize_axis_for_numpy(axis)
+    return np.array_split(x, indices_or_sections, axis=axis)
+
+
 def stack(x, axis=0):
     axis = standardize_axis_for_numpy(axis)
     dtype_set = set([getattr(a, "dtype", type(a)) for a in x])
@@ -1172,8 +1208,10 @@ def trace(x, offset=0, axis1=0, axis2=1):
     axis2 = standardize_axis_for_numpy(axis2)
     x = convert_to_tensor(x)
     dtype = standardize_dtype(x.dtype)
-    if dtype not in ("int64", "uint32", "uint64"):
-        dtype = dtypes.result_type(dtype, "int32")
+    if dtype in ("bool", "int8", "int16"):
+        dtype = "int32"
+    elif dtype in ("uint8", "uint16"):
+        dtype = "uint32"
     return np.trace(x, offset=offset, axis1=axis1, axis2=axis2, dtype=dtype)
 
 
@@ -1301,6 +1339,14 @@ def negative(x):
     return np.negative(x)
 
 
+def nextafter(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
+
+    return np.nextafter(x1, x2).astype(dtype)
+
+
 def square(x):
     x = convert_to_tensor(x)
     if standardize_dtype(x.dtype) == "bool":
@@ -1329,6 +1375,23 @@ def transpose(x, axes=None):
     return np.transpose(x, axes=axes)
 
 
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    y = convert_to_tensor(y)
+    result_dtype = dtypes.result_type(y.dtype, float)
+    if x is not None:
+        x = convert_to_tensor(x)
+    dx = convert_to_tensor(dx)
+    return np.trapezoid(y, x, dx=dx, axis=axis).astype(result_dtype)
+
+
+def vander(x, N=None, increasing=False):
+    x = convert_to_tensor(x)
+    result_dtype = dtypes.result_type(x.dtype)
+    compute_dtype = dtypes.result_type(x.dtype, config.floatx())
+    x = x.astype(compute_dtype)
+    return np.vander(x, N=N, increasing=increasing).astype(result_dtype)
+
+
 def var(x, axis=None, keepdims=False):
     axis = standardize_axis_for_numpy(axis)
     x = convert_to_tensor(x)

keras/src/backend/openvino/__init__.py

@@ -15,6 +15,7 @@ from keras.src.backend.openvino.core import compute_output_spec
 from keras.src.backend.openvino.core import cond
 from keras.src.backend.openvino.core import convert_to_numpy
 from keras.src.backend.openvino.core import convert_to_tensor
+from keras.src.backend.openvino.core import device_scope
 from keras.src.backend.openvino.core import is_tensor
 from keras.src.backend.openvino.core import random_seed_dtype
 from keras.src.backend.openvino.core import shape

keras/src/backend/openvino/core.py

@@ -13,7 +13,6 @@ from openvino import compile_model
 from keras.src import tree
 from keras.src.backend.common import KerasVariable
 from keras.src.backend.common import dtypes
-from keras.src.backend.common import global_state
 from keras.src.backend.common import standardize_dtype
 from keras.src.backend.common.dtypes import result_type
 from keras.src.backend.common.keras_tensor import KerasTensor
@@ -530,31 +529,11 @@ def ov_to_keras_type(ov_type):
 
 @contextlib.contextmanager
 def device_scope(device_name):
-    current_device = _parse_device_input(device_name)
-    global_state.set_global_attribute("openvino_device", current_device)
+    yield
 
 
 def get_device():
-    device = global_state.get_global_attribute("openvino_device", None)
-    if device is None:
-        return "CPU"
-    return device
-
-
-def _parse_device_input(device_name):
-    if isinstance(device_name, str):
-        # We support string value like "cpu:0", "gpu:1", and need to convert
-        # "gpu" to "cuda"
-        device_name = device_name.upper()
-        device_type, _ = device_name.split(":")
-        return device_type
-    else:
-        raise ValueError(
-            "Invalid value for argument `device_name`. "
-            "Expected a string like 'gpu:0' or 'cpu'. "
-            f"Received: device_name='{device_name}'"
-        )
-    return device_name
+    return "CPU"
 
 
 class Variable(KerasVariable):

keras/src/backend/openvino/linalg.py

@@ -56,3 +56,7 @@ def svd(x, full_matrices=True, compute_uv=True):
 
 def lstsq(a, b, rcond=None):
     raise NotImplementedError("`lstsq` is not supported with openvino backend")
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    raise NotImplementedError("`jvp` is not supported with openvino backend")