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

keras/src/backend/openvino/random.py

@@ -22,21 +22,14 @@ def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
 
 def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
     dtype = dtype or floatx()
-    ov_type = OPENVINO_DTYPES[dtype]
-    seed = draw_seed(seed)
-    if isinstance(seed, OpenVINOKerasTensor):
-        seed1, seed2 = convert_to_numpy(seed)
+    seed_val = draw_seed(seed)
+    if isinstance(seed_val, OpenVINOKerasTensor):
+        seed_data = convert_to_numpy(seed_val)
     else:
-        seed1, seed2 = draw_seed(seed).data
-    minval_const = ov_opset.constant(minval, dtype=dtype)
-    maxval_const = ov_opset.constant(maxval, dtype=dtype)
-    if isinstance(shape, tuple):
-        shape = list(shape)
-    output_shape_const = ov_opset.constant(shape, dtype=Type.i32)
-    random_uniform = ov_opset.random_uniform(
-        output_shape_const, minval_const, maxval_const, ov_type, seed1, seed2
-    )
-    return OpenVINOKerasTensor(random_uniform.output(0))
+        seed_data = seed_val.data
+    rng = np.random.default_rng(seed_data)
+    random_values = rng.uniform(minval, maxval, size=shape).astype(dtype)
+    return OpenVINOKerasTensor(ov_opset.constant(random_values).output(0))
 
 
 def categorical(logits, num_samples, dtype="int64", seed=None):

keras/src/backend/tensorflow/layer.py

@@ -13,7 +13,6 @@ class TFLayer(KerasAutoTrackable):
         self._saved_model_arg_spec = None
         self._tracked = []
 
-    @tf.__internal__.tracking.no_automatic_dependency_tracking
     def _set_save_spec(self, inputs, args=None, kwargs=None):
         """Defines the save spec so that serialization can trace layer calls.
 
@@ -45,6 +44,7 @@ class TFLayer(KerasAutoTrackable):
             kwargs_spec,
         )
 
+    @tf.__internal__.tracking.no_automatic_dependency_tracking
     def _trackable_children(self, save_type="checkpoint", **kwargs):
         if save_type == "savedmodel":
             # SavedModel needs to ignore the execution functions.
@@ -62,17 +62,51 @@ class TFLayer(KerasAutoTrackable):
             self.test_function = test_function
             self.predict_function = predict_function
 
-            for tracked_attr in self._tracked:
-                tracked_item = getattr(self, tracked_attr)
-                if isinstance(tracked_item, tracking.TrackedList):
-                    children[tracked_attr] = list(tracked_item)
-                if isinstance(tracked_item, tracking.TrackedDict):
-                    children[tracked_attr] = dict(tracked_item)
-                if isinstance(tracked_item, tracking.TrackedSet):
-                    children[tracked_attr] = list(tracked_item)
+            # Convert Keras tracked collections to plain Python structures
+            # without creating TensorFlow trackable dependencies
+            self._convert_tracked_collections(children)
 
         return children
 
+    def _convert_tracked_collections(self, children):
+        """Convert TrackedList/Dict/Set to plain Python structures."""
+        for tracked_attr in self._tracked:
+            tracked_item = getattr(self, tracked_attr)
+            if isinstance(tracked_item, tracking.TrackedList):
+                children[tracked_attr] = list(tracked_item)
+            if isinstance(tracked_item, tracking.TrackedDict):
+                children[tracked_attr] = dict(tracked_item)
+            if isinstance(tracked_item, tracking.TrackedSet):
+                children[tracked_attr] = list(tracked_item)
+
+    def _get_save_spec(self, dynamic_batch=True):
+        """Compatibility shim for TensorFlow saving utilities.
+
+        TensorFlow's SavedModel / TFLite export paths (e.g.,
+        tf.lite.TFLiteConverter.from_keras_model) expect a `_get_save_spec`
+        method on models. This method generates TensorSpec objects
+        describing the model's input signature.
+
+        Args:
+            dynamic_batch: whether to set the batch dimension to `None`.
+
+        Returns:
+            A TensorSpec, list or dict mirroring the model inputs, or
+            `None` when specs cannot be inferred.
+        """
+        # Lazy import to avoid circular dependency
+        from keras.src.export.export_utils import make_tf_tensor_spec
+
+        # Fall back to building specs from `self.inputs`
+        inputs = getattr(self, "inputs", None)
+        if inputs is None:
+            return None
+
+        return tree.map_structure(
+            lambda x: make_tf_tensor_spec(x, dynamic_batch=dynamic_batch),
+            inputs,
+        )
+
     @property
     def _default_save_signature(self):
         """For SavedModel support: returns the default serving signature."""

keras/src/backend/tensorflow/linalg.py

@@ -244,3 +244,27 @@ def lstsq(a, b, rcond=None):
     if b_orig_ndim == 1:
         x = tf.reshape(x, [-1])
     return x
+
+
+def jvp(fun, primals, tangents, has_aux=False):
+    primal_flat = tf.nest.flatten(primals)
+    tangent_flat = tf.nest.flatten(tangents)
+
+    tangent_flat = [
+        tf.cast(t, p.dtype) for t, p in zip(tangent_flat, primal_flat)
+    ]
+
+    with tf.autodiff.ForwardAccumulator(primal_flat, tangent_flat) as acc:
+        if has_aux:
+            primals_out, aux = fun(*primals)
+        else:
+            primals_out = fun(*primals)
+
+        primals_out_flat = tf.nest.flatten(primals_out)
+        tangents_out_flat = [acc.jvp(po) for po in primals_out_flat]
+
+    tangents_out = tf.nest.pack_sequence_as(primals_out, tangents_out_flat)
+
+    if has_aux:
+        return primals_out, tangents_out, aux
+    return primals_out, tangents_out

keras/src/backend/tensorflow/nn.py

@@ -4,6 +4,9 @@ import warnings
 import tensorflow as tf
 
 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_output_shape,
 )
@@ -268,6 +271,486 @@ def average_pool(
     return outputs
 
 
+def _compute_static_gather_indices(
+    input_dim, output_size, small_window, big_window
+):
+    """Compute gather indices for Two-Pool Gather method (corrected)."""
+    window_starts = tf.cast(
+        tf.floor(
+            tf.cast(tf.range(output_size), tf.float32)
+            * tf.cast(input_dim, tf.float32)
+            / tf.cast(output_size, tf.float32)
+        ),
+        tf.int32,
+    )
+    window_ends = tf.cast(
+        tf.math.ceil(
+            tf.cast(tf.range(1, output_size + 1), tf.float32)
+            * tf.cast(input_dim, tf.float32)
+            / tf.cast(output_size, tf.float32)
+        ),
+        tf.int32,
+    )
+
+    window_ends = tf.minimum(window_ends, input_dim)
+    window_starts = tf.minimum(window_starts, input_dim - 1)
+
+    window_sizes = window_ends - window_starts
+    is_big_window = tf.equal(window_sizes, big_window)
+
+    small_pool_len = max(1, input_dim - small_window + 1)
+
+    small_indices = window_starts
+    big_indices = window_starts + small_pool_len
+
+    gather_indices = tf.where(is_big_window, big_indices, small_indices)
+    return tf.cast(gather_indices, tf.int32)
+
+
+def _adaptive_average_pool1d(inputs, output_size, data_format="channels_first"):
+    if isinstance(output_size, int):
+        output_size = (output_size,)
+    if data_format == "channels_first":
+        inputs = tf.transpose(inputs, (0, 2, 1))
+
+    static_shape = inputs.shape.as_list()
+    l_static = static_shape[1]
+    out_l = output_size[0]
+
+    if l_static is None:
+        raise ValueError(
+            "Input length must be statically known for adaptive pooling"
+        )
+
+    small_l, big_l = compute_adaptive_pooling_window_sizes(l_static, out_l)
+    gather_l = _compute_static_gather_indices(l_static, out_l, small_l, big_l)
+
+    small_pool_l = tf.nn.pool(
+        inputs,
+        window_shape=(small_l,),
+        pooling_type="AVG",
+        strides=(1,),
+        padding="VALID",
+        data_format="NWC",
+    )
+    big_pool_l = tf.nn.pool(
+        inputs,
+        window_shape=(big_l,),
+        pooling_type="AVG",
+        strides=(1,),
+        padding="VALID",
+        data_format="NWC",
+    )
+
+    combined_l = tf.concat([small_pool_l, big_pool_l], axis=1)
+    pooled_l = tf.gather(combined_l, gather_l, axis=1)
+
+    if data_format == "channels_first":
+        pooled_l = tf.transpose(pooled_l, (0, 2, 1))
+    return pooled_l
+
+
+def _adaptive_max_pool1d(inputs, output_size, data_format="channels_first"):
+    if isinstance(output_size, int):
+        output_size = (output_size,)
+    if data_format == "channels_first":
+        inputs = tf.transpose(inputs, (0, 2, 1))
+
+    static_shape = inputs.shape.as_list()
+    l_static = static_shape[1]
+    out_l = output_size[0]
+
+    if l_static is None:
+        raise ValueError(
+            "Input length must be statically known for adaptive pooling"
+        )
+
+    small_l, big_l = compute_adaptive_pooling_window_sizes(l_static, out_l)
+    gather_l = _compute_static_gather_indices(l_static, out_l, small_l, big_l)
+
+    small_pool_l = tf.nn.pool(
+        inputs,
+        window_shape=(small_l,),
+        pooling_type="MAX",
+        strides=(1,),
+        padding="VALID",
+        data_format="NWC",
+    )
+    big_pool_l = tf.nn.pool(
+        inputs,
+        window_shape=(big_l,),
+        pooling_type="MAX",
+        strides=(1,),
+        padding="VALID",
+        data_format="NWC",
+    )
+
+    combined_l = tf.concat([small_pool_l, big_pool_l], axis=1)
+    pooled_l = tf.gather(combined_l, gather_l, axis=1)
+
+    if data_format == "channels_first":
+        pooled_l = tf.transpose(pooled_l, (0, 2, 1))
+    return pooled_l
+
+
+def _adaptive_average_pool2d(inputs, output_size, data_format="channels_first"):
+    if isinstance(output_size, int):
+        output_size = (output_size, output_size)
+
+    if data_format == "channels_first":
+        inputs = tf.transpose(inputs, (0, 2, 3, 1))
+
+    static_shape = inputs.shape.as_list()
+    h_static = static_shape[1]
+    w_static = static_shape[2]
+    out_h, out_w = output_size
+
+    if h_static is None or w_static is None:
+        raise ValueError(
+            "Input spatial dimensions must be "
+            "statically known for adaptive pooling"
+        )
+
+    small_h, big_h = compute_adaptive_pooling_window_sizes(h_static, out_h)
+    small_w, big_w = compute_adaptive_pooling_window_sizes(w_static, out_w)
+
+    gather_h = _compute_static_gather_indices(h_static, out_h, small_h, big_h)
+    gather_w = _compute_static_gather_indices(w_static, out_w, small_w, big_w)
+
+    small_pool_h = tf.nn.pool(
+        inputs,
+        window_shape=(small_h, 1),
+        pooling_type="AVG",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+    big_pool_h = tf.nn.pool(
+        inputs,
+        window_shape=(big_h, 1),
+        pooling_type="AVG",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+
+    combined_h = tf.concat([small_pool_h, big_pool_h], axis=1)
+    pooled_h = tf.gather(combined_h, gather_h, axis=1)
+
+    small_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, small_w),
+        pooling_type="AVG",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+    big_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, big_w),
+        pooling_type="AVG",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+
+    combined_w = tf.concat([small_pool_w, big_pool_w], axis=2)
+    pooled_w = tf.gather(combined_w, gather_w, axis=2)
+
+    if data_format == "channels_first":
+        pooled_w = tf.transpose(pooled_w, (0, 3, 1, 2))
+
+    return pooled_w
+
+
+def _adaptive_max_pool2d(inputs, output_size, data_format="channels_first"):
+    """Adaptive Max Pooling 2D using Two-Pool Gather method."""
+    if isinstance(output_size, int):
+        output_size = (output_size, output_size)
+
+    if data_format == "channels_first":
+        inputs = tf.transpose(inputs, (0, 2, 3, 1))
+
+    static_shape = inputs.shape.as_list()
+    h_static = static_shape[1]
+    w_static = static_shape[2]
+    out_h, out_w = output_size
+
+    if h_static is None or w_static is None:
+        raise ValueError(
+            "Input spatial dimensions must be "
+            "statically known for adaptive pooling"
+        )
+
+    small_h, big_h = compute_adaptive_pooling_window_sizes(h_static, out_h)
+    small_w, big_w = compute_adaptive_pooling_window_sizes(w_static, out_w)
+
+    gather_h = _compute_static_gather_indices(h_static, out_h, small_h, big_h)
+    gather_w = _compute_static_gather_indices(w_static, out_w, small_w, big_w)
+
+    small_pool_h = tf.nn.pool(
+        inputs,
+        window_shape=(small_h, 1),
+        pooling_type="MAX",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+    big_pool_h = tf.nn.pool(
+        inputs,
+        window_shape=(big_h, 1),
+        pooling_type="MAX",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+
+    combined_h = tf.concat([small_pool_h, big_pool_h], axis=1)
+    pooled_h = tf.gather(combined_h, gather_h, axis=1)
+
+    small_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, small_w),
+        pooling_type="MAX",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+    big_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, big_w),
+        pooling_type="MAX",
+        strides=(1, 1),
+        padding="VALID",
+        data_format="NHWC",
+    )
+
+    combined_w = tf.concat([small_pool_w, big_pool_w], axis=2)
+    pooled_w = tf.gather(combined_w, gather_w, axis=2)
+
+    if data_format == "channels_first":
+        pooled_w = tf.transpose(pooled_w, (0, 3, 1, 2))
+
+    return pooled_w
+
+
+def _adaptive_average_pool3d(inputs, output_size, data_format="channels_first"):
+    if isinstance(output_size, int):
+        output_size = (output_size, output_size, output_size)
+
+    if data_format == "channels_first":
+        inputs = tf.transpose(inputs, (0, 2, 3, 4, 1))
+
+    static_shape = inputs.shape.as_list()
+    d_static = static_shape[1]
+    h_static = static_shape[2]
+    w_static = static_shape[3]
+    out_d, out_h, out_w = output_size
+
+    if d_static is None or h_static is None or w_static is None:
+        raise ValueError(
+            "Input spatial dimensions must be "
+            "statically known for adaptive pooling"
+        )
+
+    small_d, big_d = compute_adaptive_pooling_window_sizes(d_static, out_d)
+    small_h, big_h = compute_adaptive_pooling_window_sizes(h_static, out_h)
+    small_w, big_w = compute_adaptive_pooling_window_sizes(w_static, out_w)
+
+    gather_d = _compute_static_gather_indices(d_static, out_d, small_d, big_d)
+    gather_h = _compute_static_gather_indices(h_static, out_h, small_h, big_h)
+    gather_w = _compute_static_gather_indices(w_static, out_w, small_w, big_w)
+
+    small_pool_d = tf.nn.pool(
+        inputs,
+        window_shape=(small_d, 1, 1),
+        pooling_type="AVG",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+    big_pool_d = tf.nn.pool(
+        inputs,
+        window_shape=(big_d, 1, 1),
+        pooling_type="AVG",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+
+    combined_d = tf.concat([small_pool_d, big_pool_d], axis=1)
+    pooled_d = tf.gather(combined_d, gather_d, axis=1)
+
+    small_pool_h = tf.nn.pool(
+        pooled_d,
+        window_shape=(1, small_h, 1),
+        pooling_type="AVG",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+    big_pool_h = tf.nn.pool(
+        pooled_d,
+        window_shape=(1, big_h, 1),
+        pooling_type="AVG",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+
+    combined_h = tf.concat([small_pool_h, big_pool_h], axis=2)
+    pooled_h = tf.gather(combined_h, gather_h, axis=2)
+
+    small_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, 1, small_w),
+        pooling_type="AVG",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+    big_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, 1, big_w),
+        pooling_type="AVG",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+
+    combined_w = tf.concat([small_pool_w, big_pool_w], axis=3)
+    pooled_w = tf.gather(combined_w, gather_w, axis=3)
+
+    if data_format == "channels_first":
+        pooled_w = tf.transpose(pooled_w, (0, 4, 1, 2, 3))
+
+    return pooled_w
+
+
+def _adaptive_max_pool3d(inputs, output_size, data_format="channels_first"):
+    """Adaptive Max Pooling 3D using Two-Pool Gather method."""
+    if isinstance(output_size, int):
+        output_size = (output_size, output_size, output_size)
+
+    if data_format == "channels_first":
+        inputs = tf.transpose(inputs, (0, 2, 3, 4, 1))
+
+    static_shape = inputs.shape.as_list()
+    d_static = static_shape[1]
+    h_static = static_shape[2]
+    w_static = static_shape[3]
+    out_d, out_h, out_w = output_size
+
+    if d_static is None or h_static is None or w_static is None:
+        raise ValueError(
+            "Input spatial dimensions must be "
+            "statically known for adaptive pooling"
+        )
+
+    small_d, big_d = compute_adaptive_pooling_window_sizes(d_static, out_d)
+    small_h, big_h = compute_adaptive_pooling_window_sizes(h_static, out_h)
+    small_w, big_w = compute_adaptive_pooling_window_sizes(w_static, out_w)
+
+    gather_d = _compute_static_gather_indices(d_static, out_d, small_d, big_d)
+    gather_h = _compute_static_gather_indices(h_static, out_h, small_h, big_h)
+    gather_w = _compute_static_gather_indices(w_static, out_w, small_w, big_w)
+
+    small_pool_d = tf.nn.pool(
+        inputs,
+        window_shape=(small_d, 1, 1),
+        pooling_type="MAX",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+    big_pool_d = tf.nn.pool(
+        inputs,
+        window_shape=(big_d, 1, 1),
+        pooling_type="MAX",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+
+    combined_d = tf.concat([small_pool_d, big_pool_d], axis=1)
+    pooled_d = tf.gather(combined_d, gather_d, axis=1)
+
+    small_pool_h = tf.nn.pool(
+        pooled_d,
+        window_shape=(1, small_h, 1),
+        pooling_type="MAX",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+    big_pool_h = tf.nn.pool(
+        pooled_d,
+        window_shape=(1, big_h, 1),
+        pooling_type="MAX",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+
+    combined_h = tf.concat([small_pool_h, big_pool_h], axis=2)
+    pooled_h = tf.gather(combined_h, gather_h, axis=2)
+
+    small_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, 1, small_w),
+        pooling_type="MAX",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+    big_pool_w = tf.nn.pool(
+        pooled_h,
+        window_shape=(1, 1, big_w),
+        pooling_type="MAX",
+        strides=(1, 1, 1),
+        padding="VALID",
+        data_format="NDHWC",
+    )
+
+    combined_w = tf.concat([small_pool_w, big_pool_w], axis=3)
+    pooled_w = tf.gather(combined_w, gather_w, axis=3)
+
+    if data_format == "channels_first":
+        pooled_w = tf.transpose(pooled_w, (0, 4, 1, 2, 3))
+
+    return pooled_w
+
+
+def adaptive_average_pool(inputs, output_size, data_format=None):
+    data_format = backend.standardize_data_format(data_format)
+    ndims = len(inputs.shape) - 2
+    if ndims == 1:
+        return _adaptive_average_pool1d(inputs, output_size, data_format)
+    elif ndims == 2:
+        return _adaptive_average_pool2d(inputs, output_size, data_format)
+    elif ndims == 3:
+        return _adaptive_average_pool3d(inputs, output_size, data_format)
+    else:
+        raise ValueError(
+            "adaptive_average_pool supports 1D, 2D, or 3D inputs only."
+        )
+
+
+def adaptive_max_pool(inputs, output_size, data_format=None):
+    data_format = backend.standardize_data_format(data_format)
+    ndims = len(inputs.shape) - 2
+    if ndims == 1:
+        return _adaptive_max_pool1d(inputs, output_size, data_format)
+    elif ndims == 2:
+        return _adaptive_max_pool2d(inputs, output_size, data_format)
+    elif ndims == 3:
+        return _adaptive_max_pool3d(inputs, output_size, data_format)
+    else:
+        raise ValueError(
+            "adaptive_max_pool supports 1D, 2D, or 3D inputs only."
+        )
+
+
 def _convert_data_format(data_format, ndim):
     if data_format == "channels_last":
         if ndim == 3:
@@ -310,7 +793,7 @@ def conv(
 ):
     def _conv():
         tf_data_format = _convert_data_format(data_format, len(inputs.shape))
-        return tf.nn.convolution(
+        result = tf.nn.convolution(
             inputs,
             kernel,
             strides,
@@ -318,6 +801,20 @@ def conv(
             data_format=tf_data_format,
             dilations=dilation_rate,
         )
+        result_shape = result.shape
+        if (
+            result_shape.is_fully_defined()
+            and math.prod(result_shape.as_list()) == 0
+        ):
+            raise ValueError(
+                "The convolution operation resulted in an empty output. "
+                "Output shape:"
+                f" {result_shape}. 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
 
     # Certain ops are are broken in Tensorflow on CPU only.
     # We can work around by compiling the op with XLA.
@@ -1077,3 +1574,50 @@ 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):
+    """Tensorflow 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)
+    """
+    k = (
+        (kernel_size, kernel_size)
+        if isinstance(kernel_size, int)
+        else kernel_size
+    )
+    d = (dilation, dilation) if isinstance(dilation, int) else dilation
+    p = (padding, padding) if isinstance(padding, int) else padding
+    s = (stride, stride) if isinstance(stride, int) else stride
+    N, C, H, W = input.shape
+
+    # ---- padding ----
+    if any(_ > 0 for _ in p):
+        input = tf.pad(input, [[0, 0], [0, 0], [p[0], p[0]], [p[1], p[1]]])
+    x = tf.transpose(input, [0, 2, 3, 1])  # (N, H, W, C)
+    patches = tf.image.extract_patches(
+        images=x,
+        sizes=[1, k[0], k[1], 1],
+        strides=[1, s[0], s[1], 1],
+        rates=[1, d[0], d[1], 1],
+        padding="VALID",
+    )  # (N, nH, nW, kH*kW*C)
+
+    N, nH, nW, D = patches.shape
+    patches = tf.reshape(
+        patches, [N, nH, nW, k[0], k[1], C]
+    )  # (N, nH, nW, kH, kW, C)
+    patches = tf.transpose(
+        patches, [0, 5, 3, 4, 1, 2]
+    )  # (N, C, kH, kW, nH, nW)
+    patches = tf.reshape(patches, [N, C * k[0] * k[1], nH * nW])
+    return patches