remove: items slated for removal in v0.4.0

Refs: #785

coreax/util.py

@@ -45,7 +45,7 @@ class factories and checks for numerical precision.
 import jax.tree_util as jtu
 from jax import Array, block_until_ready, jit, vmap
 from jaxtyping import Shaped
-from typing_extensions import TypeAlias, deprecated
+from typing_extensions import TypeAlias
 
 _logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO, stream=sys.stdout)
@@ -236,37 +236,6 @@ def difference(
     return x - y
 
 
-@deprecated(
-    "Use coreax.kernels.util.median_heuristic instead."
-    + " Deprecated since version 0.3.0."
-    + " Will be removed in version 0.4.0."
-)
-@jit
-def median_heuristic(
-    x: Union[Shaped[Array, " n d"], Shaped[Array, " n"], Shaped[Array, ""], float, int],
-) -> Shaped[Array, ""]:
-    """
-    Compute the median heuristic for setting kernel bandwidth.
-
-    Analysis of the performance of the median heuristic can be found in
-    :cite:`garreau2018median`.
-
-    :param x: Input array of vectors
-    :return: Bandwidth parameter, computed from the median heuristic, as a
-        zero-dimensional array
-    """
-    # Format inputs
-    x = jnp.atleast_2d(x)
-    # Calculate square distances as an upper triangular matrix
-    square_distances = jnp.triu(pairwise(squared_distance)(x, x), k=1)
-    # Calculate the median of the square distances
-    median_square_distance = jnp.median(
-        square_distances[jnp.triu_indices_from(square_distances, k=1)]
-    )
-
-    return jnp.sqrt(median_square_distance / 2.0)
-
-
 def sample_batch_indices(
     random_key: KeyArrayLike,
     max_index: int,