def __call__(
        self,
        apply_fn: Callable[
            [Dict, InputData], Union[jnp.ndarray, Tuple[jnp.ndarray, Dict]]
        ],
        model_params: Dict,
        x: Any,
        has_aux: bool,
    ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Dict]]:
        """
        Apply a transformation to the forward pass.

        Parameters
        ----------
        apply_fn: Callable[[Dict, InputData], Union[jnp.ndarray, Tuple[jnp.ndarray, Dict]]]
            The model forward pass.
        model_params: Dict
            The model parameters.
        x: Array
            Batch of inputs.
        has_aux: bool
            Whether the forward pass returns auxiliary objects.

        Returns
        -------
        Union[jnp.ndarray, Tuple[jnp.ndarray, Dict]]
            Return the transformed outputs, and auxiliary objects if available.
        """
        pass

        Apply a transformation to the forward pass.

        Parameters
        ----------
        apply_fn: Callable[[Dict, InputData], Union[jnp.ndarray, Tuple[jnp.ndarray, Dict]]]
            The model forward pass.
        model_params: Dict
            The model parameters.
        x: Array
            Batch of inputs.
        has_aux: bool
            Whether the forward pass returns auxiliary objects.

        Returns
        -------
        Union[jnp.ndarray, Tuple[jnp.ndarray, Dict]]
            Return the transformed outputs, and auxiliary objects if available.
        
fortuna/model_editor/base.py