"""
Colorimetry Plotting
====================

Defines the colorimetry plotting objects:

-   :func:`colour.plotting.plot_single_sd`
-   :func:`colour.plotting.plot_multi_sds`
-   :func:`colour.plotting.plot_single_cmfs`
-   :func:`colour.plotting.plot_multi_cmfs`
-   :func:`colour.plotting.plot_single_illuminant_sd`
-   :func:`colour.plotting.plot_multi_illuminant_sds`
-   :func:`colour.plotting.plot_visible_spectrum`
-   :func:`colour.plotting.plot_single_lightness_function`
-   :func:`colour.plotting.plot_multi_lightness_functions`
-   :func:`colour.plotting.plot_single_luminance_function`
-   :func:`colour.plotting.plot_multi_luminance_functions`
-   :func:`colour.plotting.plot_blackbody_spectral_radiance`
-   :func:`colour.plotting.plot_blackbody_colours`

References
----------
-   :cite:`Spiker2015a` : Borer, T. (2017). Private Discussion with Mansencal,
    T. and Shaw, N.
"""

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np
from functools import reduce
from matplotlib.patches import Polygon

from colour.algebra import (
    LinearInterpolator,
    normalise_maximum,
    sdiv,
    sdiv_mode,
)
from colour.colorimetry import (
    CCS_ILLUMINANTS,
    SDS_ILLUMINANTS,
    LIGHTNESS_METHODS,
    LUMINANCE_METHODS,
    MultiSpectralDistributions,
    SpectralDistribution,
    SpectralShape,
    sd_blackbody,
    sd_ones,
    sd_to_XYZ,
    sds_and_msds_to_sds,
    wavelength_to_XYZ,
)
from colour.hints import (
    Any,
    Boolean,
    Callable,
    Dict,
    Floating,
    List,
    NDArray,
    Optional,
    Sequence,
    Tuple,
    Union,
    cast,
)
from colour.plotting import (
    CONSTANTS_COLOUR_STYLE,
    XYZ_to_plotting_colourspace,
    artist,
    filter_passthrough,
    filter_cmfs,
    filter_illuminants,
    override_style,
    render,
    plot_single_colour_swatch,
    plot_multi_functions,
    update_settings_collection,
)
from colour.utilities import (
    as_float_array,
    domain_range_scale,
    first_item,
    ones,
    tstack,
)

__author__ = "Colour Developers"
__copyright__ = "Copyright 2013 Colour Developers"
__license__ = "New BSD License - https://opensource.org/licenses/BSD-3-Clause"
__maintainer__ = "Colour Developers"
__email__ = "colour-developers@colour-science.org"
__status__ = "Production"

__all__ = [
    "plot_single_sd",
    "plot_multi_sds",
    "plot_single_cmfs",
    "plot_multi_cmfs",
    "plot_single_illuminant_sd",
    "plot_multi_illuminant_sds",
    "plot_visible_spectrum",
    "plot_single_lightness_function",
    "plot_multi_lightness_functions",
    "plot_single_luminance_function",
    "plot_multi_luminance_functions",
    "plot_blackbody_spectral_radiance",
    "plot_blackbody_colours",
]

@override_style()
def plot_single_sd(
    sd: SpectralDistribution,
    cmfs: Union[
        MultiSpectralDistributions,
        str,
        Sequence[Union[MultiSpectralDistributions, str]],
    ] = "CIE 1931 2 Degree Standard Observer",
    out_of_gamut_clipping: Boolean = True,
    modulate_colours_with_sd_amplitude: Boolean = False,
    equalize_sd_amplitude: Boolean = False,
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given spectral distribution.

    Parameters
    ----------
    sd
        Spectral distribution to plot.
    cmfs
        Standard observer colour matching functions used for computing the
        spectrum domain and colours. ``cmfs`` can be of any type or form
        supported by the :func:`colour.plotting.filter_cmfs` definition.
    out_of_gamut_clipping
        Whether to clip out of gamut colours otherwise, the colours will be
        offset by the absolute minimal colour leading to a rendering on
        gray background, less saturated and smoother.
    modulate_colours_with_sd_amplitude
        Whether to modulate the colours with the spectral distribution
        amplitude.
    equalize_sd_amplitude
        Whether to equalize the spectral distribution amplitude.
        Equalization occurs after the colours modulation thus setting both
        arguments to *True* will generate a spectrum strip where each
        wavelength colour is modulated by the spectral distribution amplitude.
        The usual 5% margin above the spectral distribution is also omitted.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`, :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    References
    ----------
    :cite:`Spiker2015a`

    Examples
    --------
    >>> from colour import SpectralDistribution
    >>> data = {
    ...     500: 0.0651,
    ...     520: 0.0705,
    ...     540: 0.0772,
    ...     560: 0.0870,
    ...     580: 0.1128,
    ...     600: 0.1360
    ... }
    >>> sd = SpectralDistribution(data, name='Custom')
    >>> plot_single_sd(sd)  # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Single_SD.png
        :align: center
        :alt: plot_single_sd
    """

    _figure, axes = artist(**kwargs)

    cmfs = cast(
        MultiSpectralDistributions, first_item(filter_cmfs(cmfs).values())
    )

    sd = cast(SpectralDistribution, sd.copy())
    sd.interpolator = LinearInterpolator
    wavelengths = cmfs.wavelengths[
        np.logical_and(
            cmfs.wavelengths
            >= max(min(cmfs.wavelengths), min(sd.wavelengths)),
            cmfs.wavelengths
            <= min(max(cmfs.wavelengths), max(sd.wavelengths)),
        )
    ]
    values = as_float_array(sd[wavelengths])

    RGB = XYZ_to_plotting_colourspace(
        wavelength_to_XYZ(wavelengths, cmfs),
        CCS_ILLUMINANTS["CIE 1931 2 Degree Standard Observer"]["E"],
        apply_cctf_encoding=False,
    )

    if not out_of_gamut_clipping:
        RGB += np.abs(np.min(RGB))

    RGB = normalise_maximum(RGB)

    if modulate_colours_with_sd_amplitude:
        with sdiv_mode():
            RGB *= cast(NDArray, sdiv(values, np.max(values)))[..., np.newaxis]

    RGB = CONSTANTS_COLOUR_STYLE.colour.colourspace.cctf_encoding(RGB)

    if equalize_sd_amplitude:
        values = ones(values.shape)

    margin = 0 if equalize_sd_amplitude else 0.05

    x_min, x_max = min(wavelengths), max(wavelengths)
    y_min, y_max = 0, max(values) + max(values) * margin

    polygon = Polygon(
        np.vstack(
            [
                (x_min, 0),
                tstack([wavelengths, values]),
                (x_max, 0),
            ]
        ),
        facecolor="none",
        edgecolor="none",
        zorder=CONSTANTS_COLOUR_STYLE.zorder.background_polygon,
    )
    axes.add_patch(polygon)

    padding = 0.1
    axes.bar(
        x=wavelengths - padding,
        height=max(values),
        width=1 + padding,
        color=RGB,
        align="edge",
        clip_path=polygon,
        zorder=CONSTANTS_COLOUR_STYLE.zorder.background_polygon,
    )

    axes.plot(
        wavelengths,
        values,
        color=CONSTANTS_COLOUR_STYLE.colour.dark,
        zorder=CONSTANTS_COLOUR_STYLE.zorder.midground_line,
    )

    settings: Dict[str, Any] = {
        "axes": axes,
        "bounding_box": (x_min, x_max, y_min, y_max),
        "title": f"{sd.strict_name} - {cmfs.strict_name}",
        "x_label": "Wavelength $\\lambda$ (nm)",
        "y_label": "Spectral Distribution",
    }
    settings.update(kwargs)

    return render(**settings)

@override_style()
def plot_multi_sds(
    sds: Union[
        Sequence[Union[SpectralDistribution, MultiSpectralDistributions]],
        MultiSpectralDistributions,
    ],
    plot_kwargs: Optional[Union[Dict, List[Dict]]] = None,
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given spectral distributions.

    Parameters
    ----------
    sds
        Spectral distributions or multi-spectral distributions to
        plot. `sds` can be a single
        :class:`colour.MultiSpectralDistributions` class instance, a list
        of :class:`colour.MultiSpectralDistributions` class instances or a
        list of :class:`colour.SpectralDistribution` class instances.
    plot_kwargs
        Keyword arguments for the :func:`matplotlib.pyplot.plot` definition,
        used to control the style of the plotted spectral distributions.
        `plot_kwargs`` can be either a single dictionary applied to all the
        plotted spectral distributions with the same settings or a sequence of
        dictionaries with different settings for each plotted spectral
        distributions. The following special keyword arguments can also be
        used:

        -   ``illuminant`` : The illuminant used to compute the spectral
            distributions colours. The default is the illuminant associated
            with the whitepoint of the default plotting colourspace.
            ``illuminant`` can be of any type or form supported by the
            :func:`colour.plotting.filter_cmfs` definition.
        -   ``cmfs`` : The standard observer colour matching functions used for
            computing the spectral distributions colours. ``cmfs`` can be of
            any type or form supported by the
            :func:`colour.plotting.filter_cmfs` definition.
        -   ``normalise_sd_colours`` : Whether to normalise the computed
            spectral distributions colours. The default is *True*.
        -   ``use_sd_colours`` : Whether to use the computed spectral
            distributions colours under the plotting colourspace illuminant.
            Alternatively, it is possible to use the
            :func:`matplotlib.pyplot.plot` definition ``color`` argument with
            pre-computed values. The default is *True*.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`, :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> from colour import SpectralDistribution
    >>> data_1 = {
    ...     500: 0.004900,
    ...     510: 0.009300,
    ...     520: 0.063270,
    ...     530: 0.165500,
    ...     540: 0.290400,
    ...     550: 0.433450,
    ...     560: 0.594500
    ... }
    >>> data_2 = {
    ...     500: 0.323000,
    ...     510: 0.503000,
    ...     520: 0.710000,
    ...     530: 0.862000,
    ...     540: 0.954000,
    ...     550: 0.994950,
    ...     560: 0.995000
    ... }
    >>> sd_1 = SpectralDistribution(data_1, name='Custom 1')
    >>> sd_2 = SpectralDistribution(data_2, name='Custom 2')
    >>> plot_kwargs = [
    ...     {'use_sd_colours': True},
    ...     {'use_sd_colours': True, 'linestyle': 'dashed'},
    ... ]
    >>> plot_multi_sds([sd_1, sd_2], plot_kwargs=plot_kwargs)
    ... # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Multi_SDS.png
        :align: center
        :alt: plot_multi_sds
    """

    _figure, axes = artist(**kwargs)

    sds_converted = sds_and_msds_to_sds(sds)

    plot_settings_collection = [
        {
            "label": f"{sd.strict_name}",
            "zorder": CONSTANTS_COLOUR_STYLE.zorder.midground_line,
            "cmfs": "CIE 1931 2 Degree Standard Observer",
            "illuminant": SDS_ILLUMINANTS[
                CONSTANTS_COLOUR_STYLE.colour.colourspace.whitepoint_name
            ],
            "use_sd_colours": False,
            "normalise_sd_colours": False,
        }
        for sd in sds_converted
    ]

    if plot_kwargs is not None:
        update_settings_collection(
            plot_settings_collection, plot_kwargs, len(sds_converted)
        )

    x_limit_min, x_limit_max, y_limit_min, y_limit_max = [], [], [], []
    for i, sd in enumerate(sds_converted):
        plot_settings = plot_settings_collection[i]

        cmfs = cast(
            MultiSpectralDistributions,
            first_item(filter_cmfs(plot_settings.pop("cmfs")).values()),
        )
        illuminant = cast(
            SpectralDistribution,
            first_item(
                filter_illuminants(plot_settings.pop("illuminant")).values()
            ),
        )
        normalise_sd_colours = plot_settings.pop("normalise_sd_colours")
        use_sd_colours = plot_settings.pop("use_sd_colours")

        wavelengths, values = sd.wavelengths, sd.values

        shape = sd.shape
        x_limit_min.append(shape.start)
        x_limit_max.append(shape.end)
        y_limit_min.append(min(values))
        y_limit_max.append(max(values))

        if use_sd_colours:
            with domain_range_scale("1"):
                XYZ = sd_to_XYZ(sd, cmfs, illuminant)

            if normalise_sd_colours:
                XYZ /= XYZ[..., 1]

            plot_settings["color"] = np.clip(
                XYZ_to_plotting_colourspace(XYZ), 0, 1
            )

        axes.plot(wavelengths, values, **plot_settings)

    bounding_box = (
        min(x_limit_min),
        max(x_limit_max),
        min(y_limit_min),
        max(y_limit_max) + np.max(y_limit_max) * 0.05,
    )
    settings: Dict[str, Any] = {
        "axes": axes,
        "bounding_box": bounding_box,
        "legend": True,
        "x_label": "Wavelength $\\lambda$ (nm)",
        "y_label": "Spectral Distribution",
    }
    settings.update(kwargs)

    return render(**settings)

@override_style()
def plot_single_cmfs(
    cmfs: Union[
        MultiSpectralDistributions,
        str,
        Sequence[Union[MultiSpectralDistributions, str]],
    ] = "CIE 1931 2 Degree Standard Observer",
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given colour matching functions.

    Parameters
    ----------
    cmfs
        Colour matching functions to plot. ``cmfs`` can be of any type or form
        supported by the :func:`colour.plotting.filter_cmfs` definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_multi_cmfs`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_single_cmfs('CIE 1931 2 Degree Standard Observer')
    ... # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Single_CMFS.png
        :align: center
        :alt: plot_single_cmfs
    """

    cmfs = cast(
        MultiSpectralDistributions, first_item(filter_cmfs(cmfs).values())
    )

    settings: Dict[str, Any] = {
        "title": f"{cmfs.strict_name} - Colour Matching Functions"
    }
    settings.update(kwargs)

    return plot_multi_cmfs((cmfs,), **settings)

@override_style()
def plot_multi_cmfs(
    cmfs: Union[
        MultiSpectralDistributions,
        str,
        Sequence[Union[MultiSpectralDistributions, str]],
    ],
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given colour matching functions.

    Parameters
    ----------
    cmfs
        Colour matching functions to plot. ``cmfs`` elements can be of any
        type or form supported by the :func:`colour.plotting.filter_cmfs`
        definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`, :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> cmfs = [
    ...     'CIE 1931 2 Degree Standard Observer',
    ...     'CIE 1964 10 Degree Standard Observer',
    ... ]
    >>> plot_multi_cmfs(cmfs)  # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Multi_CMFS.png
        :align: center
        :alt: plot_multi_cmfs
    """

    cmfs = cast(
        List[MultiSpectralDistributions], list(filter_cmfs(cmfs).values())
    )

    _figure, axes = artist(**kwargs)

    axes.axhline(
        color=CONSTANTS_COLOUR_STYLE.colour.dark,
        linestyle="--",
        zorder=CONSTANTS_COLOUR_STYLE.zorder.foreground_line,
    )

    x_limit_min, x_limit_max, y_limit_min, y_limit_max = [], [], [], []
    for i, cmfs_i in enumerate(cmfs):
        for j, RGB in enumerate(
            as_float_array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
        ):
            RGB = [reduce(lambda y, _: y * 0.5, range(i), x) for x in RGB]
            values = cmfs_i.values[:, j]

            shape = cmfs_i.shape
            x_limit_min.append(shape.start)
            x_limit_max.append(shape.end)
            y_limit_min.append(np.min(values))
            y_limit_max.append(np.max(values))

            axes.plot(
                cmfs_i.wavelengths,
                values,
                color=RGB,
                label=f"{cmfs_i.strict_labels[j]} - {cmfs_i.strict_name}",
                zorder=CONSTANTS_COLOUR_STYLE.zorder.midground_line,
            )

    bounding_box = (
        min(x_limit_min),
        max(x_limit_max),
        min(y_limit_min) - np.abs(np.min(y_limit_min)) * 0.05,
        max(y_limit_max) + np.abs(np.max(y_limit_max)) * 0.05,
    )
    cmfs_strict_names = ", ".join([cmfs_i.strict_name for cmfs_i in cmfs])
    title = f"{cmfs_strict_names} - Colour Matching Functions"

    settings: Dict[str, Any] = {
        "axes": axes,
        "bounding_box": bounding_box,
        "legend": True,
        "title": title,
        "x_label": "Wavelength $\\lambda$ (nm)",
        "y_label": "Tristimulus Values",
    }
    settings.update(kwargs)

    return render(**settings)

@override_style()
def plot_single_illuminant_sd(
    illuminant: Union[SpectralDistribution, str],
    cmfs: Union[
        MultiSpectralDistributions,
        str,
        Sequence[Union[MultiSpectralDistributions, str]],
    ] = "CIE 1931 2 Degree Standard Observer",
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given single illuminant spectral distribution.

    Parameters
    ----------
    illuminant
        Illuminant to plot. ``illuminant`` can be of any type or form supported
        by the :func:`colour.plotting.filter_illuminants` definition.
    cmfs
        Standard observer colour matching functions used for computing the
        spectrum domain and colours. ``cmfs`` can be of any type or form
        supported by the :func:`colour.plotting.filter_cmfs` definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_single_sd`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    References
    ----------
    :cite:`Spiker2015a`

    Examples
    --------
    >>> plot_single_illuminant_sd('A')  # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Single_Illuminant_SD.png
        :align: center
        :alt: plot_single_illuminant_sd
    """

    cmfs = cast(
        MultiSpectralDistributions, first_item(filter_cmfs(cmfs).values())
    )

    title = f"Illuminant {illuminant} - {cmfs.strict_name}"

    illuminant = first_item(filter_illuminants(illuminant).values())

    settings: Dict[str, Any] = {"title": title, "y_label": "Relative Power"}
    settings.update(kwargs)

    return plot_single_sd(illuminant, **settings)

@override_style()
def plot_multi_illuminant_sds(
    illuminants: Union[
        SpectralDistribution, str, Sequence[Union[SpectralDistribution, str]]
    ],
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given illuminants spectral distributions.

    Parameters
    ----------
    illuminants
        Illuminants to plot. ``illuminants`` elements can be of any type or
        form supported by the :func:`colour.plotting.filter_illuminants`
        definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_multi_sds`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_multi_illuminant_sds(['A', 'B', 'C'])  # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Multi_Illuminant_SDS.png
        :align: center
        :alt: plot_multi_illuminant_sds
    """

    if "plot_kwargs" not in kwargs:
        kwargs["plot_kwargs"] = {}

    SD_E = SDS_ILLUMINANTS["E"]
    if isinstance(kwargs["plot_kwargs"], dict):
        kwargs["plot_kwargs"]["illuminant"] = SD_E
    else:
        for i in range(len(kwargs["plot_kwargs"])):
            kwargs["plot_kwargs"][i]["illuminant"] = SD_E

    illuminants = cast(
        List[SpectralDistribution],
        list(filter_illuminants(illuminants).values()),
    )

    illuminant_strict_names = ", ".join(
        [illuminant.strict_name for illuminant in illuminants]
    )
    title = f"{illuminant_strict_names} - Illuminants Spectral Distributions"

    settings: Dict[str, Any] = {"title": title, "y_label": "Relative Power"}
    settings.update(kwargs)

    return plot_multi_sds(illuminants, **settings)

@override_style(
    **{
        "ytick.left": False,
        "ytick.labelleft": False,
    }
)
def plot_visible_spectrum(
    cmfs: Union[
        MultiSpectralDistributions,
        str,
        Sequence[Union[MultiSpectralDistributions, str]],
    ] = "CIE 1931 2 Degree Standard Observer",
    out_of_gamut_clipping: Boolean = True,
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot the visible colours spectrum using given standard observer *CIE XYZ*
    colour matching functions.

    Parameters
    ----------
    cmfs
        Standard observer colour matching functions used for computing the
        spectrum domain and colours. ``cmfs`` can be of any type or form
        supported by the :func:`colour.plotting.filter_cmfs` definition.
    out_of_gamut_clipping
        Whether to clip out of gamut colours otherwise, the colours will be
        offset by the absolute minimal colour leading to a rendering on
        gray background, less saturated and smoother.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_single_sd`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    References
    ----------
    :cite:`Spiker2015a`

    Examples
    --------
    >>> plot_visible_spectrum()  # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Visible_Spectrum.png
        :align: center
        :alt: plot_visible_spectrum
    """

    cmfs = cast(
        MultiSpectralDistributions, first_item(filter_cmfs(cmfs).values())
    )

    bounding_box = (min(cmfs.wavelengths), max(cmfs.wavelengths), 0, 1)

    settings: Dict[str, Any] = {"bounding_box": bounding_box, "y_label": None}
    settings.update(kwargs)
    settings["standalone"] = False

    _figure, axes = plot_single_sd(
        sd_ones(cmfs.shape),
        cmfs=cmfs,
        out_of_gamut_clipping=out_of_gamut_clipping,
        **settings,
    )

    # Removing wavelength line as it doubles with the axes spine.
    axes.lines.pop(0)

    settings = {
        "axes": axes,
        "standalone": True,
        "title": f"The Visible Spectrum - {cmfs.strict_name}",
        "x_label": "Wavelength $\\lambda$ (nm)",
    }
    settings.update(kwargs)

    return render(**settings)

@override_style()
def plot_single_lightness_function(
    function: Union[Callable, str], **kwargs: Any
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given *Lightness* function.

    Parameters
    ----------
    function
        *Lightness* function to plot. ``function`` can be of any type or form
        supported by the :func:`colour.plotting.filter_passthrough` definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_multi_functions`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_single_lightness_function('CIE 1976')  # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Single_Lightness_Function.png
        :align: center
        :alt: plot_single_lightness_function
    """

    settings: Dict[str, Any] = {"title": f"{function} - Lightness Function"}
    settings.update(kwargs)

    return plot_multi_lightness_functions((function,), **settings)

@override_style()
def plot_multi_lightness_functions(
    functions: Union[Callable, str, Sequence[Union[Callable, str]]],
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given *Lightness* functions.

    Parameters
    ----------
    functions
        *Lightness* functions to plot. ``functions`` elements can be of any
        type or form supported by the
        :func:`colour.plotting.filter_passthrough` definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_multi_functions`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_multi_lightness_functions(['CIE 1976', 'Wyszecki 1963'])
    ... # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Multi_Lightness_Functions.png
        :align: center
        :alt: plot_multi_lightness_functions
    """

    functions_filtered = filter_passthrough(LIGHTNESS_METHODS, functions)

    settings: Dict[str, Any] = {
        "bounding_box": (0, 1, 0, 1),
        "legend": True,
        "title": f"{', '.join(functions_filtered)} - Lightness Functions",
        "x_label": "Normalised Relative Luminance Y",
        "y_label": "Normalised Lightness",
    }
    settings.update(kwargs)

    with domain_range_scale("1"):
        return plot_multi_functions(functions_filtered, **settings)

@override_style()
def plot_single_luminance_function(
    function: Union[Callable, str], **kwargs: Any
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given *Luminance* function.

    Parameters
    ----------
    function
        *Luminance* function to plot.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_multi_functions`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_single_luminance_function('CIE 1976')  # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Single_Luminance_Function.png
        :align: center
        :alt: plot_single_luminance_function
    """

    settings: Dict[str, Any] = {"title": f"{function} - Luminance Function"}
    settings.update(kwargs)

    return plot_multi_luminance_functions((function,), **settings)

@override_style()
def plot_multi_luminance_functions(
    functions: Union[Callable, str, Sequence[Union[Callable, str]]],
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given *Luminance* functions.

    Parameters
    ----------
    functions
        *Luminance* functions to plot. ``functions`` elements can be of any
        type or form supported by the
        :func:`colour.plotting.filter_passthrough` definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_multi_functions`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_multi_luminance_functions(['CIE 1976', 'Newhall 1943'])
    ... # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Multi_Luminance_Functions.png
        :align: center
        :alt: plot_multi_luminance_functions
    """

    functions_filtered = filter_passthrough(LUMINANCE_METHODS, functions)

    settings: Dict[str, Any] = {
        "bounding_box": (0, 1, 0, 1),
        "legend": True,
        "title": f"{', '.join(functions_filtered)} - Luminance Functions",
        "x_label": "Normalised Munsell Value / Lightness",
        "y_label": "Normalised Relative Luminance Y",
    }
    settings.update(kwargs)

    with domain_range_scale("1"):
        return plot_multi_functions(functions_filtered, **settings)

@override_style()
def plot_blackbody_spectral_radiance(
    temperature: Floating = 3500,
    cmfs: Union[
        MultiSpectralDistributions,
        str,
        Sequence[Union[MultiSpectralDistributions, str]],
    ] = "CIE 1931 2 Degree Standard Observer",
    blackbody: str = "VY Canis Major",
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot given blackbody spectral radiance.

    Parameters
    ----------
    temperature
        Blackbody temperature.
    cmfs
        Standard observer colour matching functions used for computing the
        spectrum domain and colours. ``cmfs`` can be of any type or form
        supported by the :func:`colour.plotting.filter_cmfs` definition.
    blackbody
        Blackbody name.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`,
        :func:`colour.plotting.plot_single_sd`,
        :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_blackbody_spectral_radiance(3500, blackbody='VY Canis Major')
    ... # doctest: +ELLIPSIS
    (<Figure size ... with 2 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Blackbody_Spectral_Radiance.png
        :align: center
        :alt: plot_blackbody_spectral_radiance
    """

    figure = plt.figure()

    figure.subplots_adjust(hspace=CONSTANTS_COLOUR_STYLE.geometry.short / 2)

    cmfs = cast(
        MultiSpectralDistributions, first_item(filter_cmfs(cmfs).values())
    )

    sd = sd_blackbody(temperature, cmfs.shape)

    axes = figure.add_subplot(211)
    settings: Dict[str, Any] = {
        "axes": axes,
        "title": f"{blackbody} - Spectral Radiance",
        "y_label": "W / (sr m$^2$) / m",
    }
    settings.update(kwargs)
    settings["standalone"] = False

    plot_single_sd(sd, cmfs.name, **settings)

    axes = figure.add_subplot(212)

    with domain_range_scale("1"):
        XYZ = sd_to_XYZ(sd, cmfs)

    RGB = normalise_maximum(XYZ_to_plotting_colourspace(XYZ))

    settings = {
        "axes": axes,
        "aspect": None,
        "title": f"{blackbody} - Colour",
        "x_label": f"{temperature}K",
        "y_label": "",
        "x_ticker": False,
        "y_ticker": False,
    }
    settings.update(kwargs)
    settings["standalone"] = False

    figure, axes = plot_single_colour_swatch(RGB, **settings)

    settings = {"axes": axes, "standalone": True}
    settings.update(kwargs)

    return render(**settings)

@override_style(
    **{
        "ytick.left": False,
        "ytick.labelleft": False,
    }
)
def plot_blackbody_colours(
    shape: SpectralShape = SpectralShape(150, 12500, 50),
    cmfs: Union[
        MultiSpectralDistributions,
        str,
        Sequence[Union[MultiSpectralDistributions, str]],
    ] = "CIE 1931 2 Degree Standard Observer",
    **kwargs: Any,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot blackbody colours.

    Parameters
    ----------
    shape
        Spectral shape to use as plot boundaries.
    cmfs
        Standard observer colour matching functions used for computing the
        blackbody colours. ``cmfs`` can be of any type or form supported by the
        :func:`colour.plotting.filter_cmfs` definition.

    Other Parameters
    ----------------
    kwargs
        {:func:`colour.plotting.artist`, :func:`colour.plotting.render`},
        See the documentation of the previously listed definitions.

    Returns
    -------
    :class:`tuple`
        Current figure and axes.

    Examples
    --------
    >>> plot_blackbody_colours(SpectralShape(150, 12500, 50))
    ... # doctest: +ELLIPSIS
    (<Figure size ... with 1 Axes>, <...AxesSubplot...>)

    .. image:: ../_static/Plotting_Plot_Blackbody_Colours.png
        :align: center
        :alt: plot_blackbody_colours
    """

    _figure, axes = artist(**kwargs)

    cmfs = cast(
        MultiSpectralDistributions, first_item(filter_cmfs(cmfs).values())
    )

    RGB = []
    temperatures = []

    for temperature in shape:
        sd = sd_blackbody(temperature, cmfs.shape)

        with domain_range_scale("1"):
            XYZ = sd_to_XYZ(sd, cmfs)

        RGB.append(normalise_maximum(XYZ_to_plotting_colourspace(XYZ)))
        temperatures.append(temperature)

    x_min, x_max = min(temperatures), max(temperatures)
    y_min, y_max = 0, 1

    padding = 0.1
    axes.bar(
        x=np.array(temperatures) - padding,
        height=1,
        width=shape.interval + (padding * shape.interval),
        color=RGB,
        align="edge",
        zorder=CONSTANTS_COLOUR_STYLE.zorder.background_polygon,
    )

    settings: Dict[str, Any] = {
        "axes": axes,
        "bounding_box": (x_min, x_max, y_min, y_max),
        "title": "Blackbody Colours",
        "x_label": "Temperature K",
        "y_label": None,
    }
    settings.update(kwargs)

    return render(**settings)


# This file is part of LayerModel_lib
#
#     A tool to compute the transmission behaviour of plane electromagnetic waves
#     through human tissue.
#
# Copyright (C) 2018 Jan-Christoph Brumm
#
# Licensed under MIT license.
#
import numpy as np
import os

from LayerModel_lib.voxelmodel import VoxelModel
from LayerModel_lib.voxelmodel_importer import VoxelModelImporter
from LayerModel_lib.coordinate import Coordinate

current_directory = os.path.dirname(__file__)
base_path = os.path.join(current_directory, '..', '..', '..', 'Numerical Human Phantoms', 'Donna')

# path to the AVW File of this model
filename = os.path.join(base_path, 'segm_donna')
# path to the tissue_mapping file
tissue_file = os.path.join('ImportDonna_tissues.txt')

AVW_Data = VoxelModelImporter(filename, tissue_file, 'AVW')
model_orig = AVW_Data.data['image']
tissue_name_orig = AVW_Data.tissue_names
tissue_mapping = AVW_Data.tissue_mapping

Donna = VoxelModel()
Donna.show_progress_bar = True

# needs to be set manually from README.txt
Donna.set_scale(1.875, 1.875, 10)

Donna.name = 'Donna'
Donna.description = 'Donna model from the Helmholtz Zentrum München. ' \
                    'Resolution %.2fmm x %.2fmm x %.2fmm' % (Donna.scaling.x,
                                                             Donna.scaling.y, Donna.scaling.z)

#  Calculate the outer_shape of the original and the complete model
outer_shape = AVW_Data.calculate_outer_shape(model_orig, tissue_mapping)

Donna.add_voxel_data(short_name='original',
                     name='Original data from AVW file',
                     model=model_orig,
                     outer_shape=outer_shape,
                     tissue_names=tissue_name_orig)

Donna.add_voxel_data(short_name='complete',
                     name='The \'original\' model converted to our TissueProperties.',
                     model=Donna.models['original'].data,
                     outer_shape=outer_shape,
                     tissue_mapping=tissue_mapping)

# Calculate the trunk model
start_slice = int(96)
end_slice = int(141)

(model_trunk, trunk_mask) = AVW_Data.calculate_trunk_model(Donna, 'complete', z_start=start_slice, z_end=end_slice)
outer_shape_trunk = AVW_Data.calculate_outer_shape(model_trunk)

Donna.add_voxel_data(short_name='trunk',
                     name="The trunk of the 'complete' model. Arms have been removed using "
                          "VoxelModel.remove_arms().",
                     outer_shape=outer_shape_trunk,
                     model=model_trunk,
                     mask=trunk_mask,
                     tissue_mapping=None)

# Calculate the 3D surface of Donna
surface = Donna.create_3d_model(model_type='trunk', patch_size=(30, 30))
Donna.models['trunk'].surface_3d = surface

Donna.models['trunk'].endpoints = []
for (i, s) in enumerate(surface):
   Donna.models['trunk'].endpoints.append(Coordinate(np.array(s['centroid'])))#

Donna.save_model()


import numpy as np
import torch

class coord_utils:
    @staticmethod
    def location_to_probability_map(size, loc):
        # loc is not normalized location
        promap_vec = torch.zeros([size, size], dtype=torch.float32)
        try:
            conf = loc[4]
        except IndexError:
            conf = 1.0

        cx = loc[0]*size
        cy = loc[1]*size
        w = loc[2]*size
        h = loc[3]*size

        [x1, y1, x2, y2] = [(cx - w/2.).int(), (cy - h/2.).int(), (cx + w/2.).int(), (cy + h/2.).int()]
        if x1 == x2: x2 += 1
        if y1 == y2: y2 += 1

        x1 = x1.clamp(0, size)
        y1 = y1.clamp(0, size)
        x2 = x2.clamp(0, size)
        y2 = y2.clamp(0, size)

        for y in range(y1, y2):
            for x in range(x1, x2):
                promap_vec[y][x] = conf
        return promap_vec

    @staticmethod
    def locations_to_probability_maps(size, locs):
        pms = []
        for loc in locs:
            pm = coord_utils.location_to_probability_map(size, loc)
            pms.append(pm.view(-1))

        return torch.stack(pms, dim=0)

    @staticmethod
    def probability_map_to_location(size, pmap):
        # probability map to location (cx, cy, w, h)
        pmap = pmap.view(size, size)

        xlist = []
        ylist = []
        for y in range(size):
            for x in range(size):
                if(pmap[y][x] >= 0.5):
                    xlist.append(x+0.5)
                    ylist.append(y+0.5)

        if len(xlist) == 0 or len(ylist) == 0:
            return torch.zeros(4, dtype=torch.float32)

        ax = np.array(xlist)
        ay = np.array(ylist)
        x1 = ax.mean()
        y1 = ay.mean()

        k = 3.5 #np.sqrt(2)
        w = ax.std() * k + 0.5
        h = ay.std() * k + 0.5

        loc = torch.tensor([x1/size, y1/size, w/size, h/size], dtype=torch.float32)
        return loc

    @staticmethod
    def normal_to_location(wid, ht, location):
        # Normalized location to coordinate
        wid *= 1.0
        ht *= 1.0
        location[0] *= wid
        location[1] *= ht
        location[2] *= wid
        location[3] *= ht
        return location

    @staticmethod
    def location_to_normal(wid, ht, location):
        # Coordinates to normalized location
        wid *= 1.0
        ht *= 1.0
        location[0] /= wid
        location[1] /= ht
        location[2] /= wid
        location[3] /= ht
        return location

    @staticmethod
    def bbox_iou(box1, box2, x1y1x2y2=True):
        """
        Returns the IoU of two bounding boxes
        """
        if not x1y1x2y2: # (cx, cy, w, h)
            # Transform from center and width to exact coordinates
            b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
            b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
            b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
            b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
        else:
            # Get the coordinates of bounding boxes
            b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
            b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

        # get the corrdinates of the intersection rectangle
        inter_rect_x1 = torch.max(b1_x1, b2_x1)
        inter_rect_y1 = torch.max(b1_y1, b2_y1)
        inter_rect_x2 = torch.min(b1_x2, b2_x2)
        inter_rect_y2 = torch.min(b1_y2, b2_y2)
        # Intersection area
        inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(
            inter_rect_y2 - inter_rect_y1 + 1, min=0
        )
        # Union Area
        b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
        b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

        iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

        return iou


import math

from typing import Optional, List, Tuple, Sequence, Union, cast, TypeVar
from typing import Iterator, overload
import numpy
import itertools

from ..types import Xp, Shape, DTypes, DTypesInt, DTypesFloat, List2d, ArrayXd
from ..types import Floats1d, Floats2d, Floats3d, Floats4d
from ..types import Array1d, Array2d, Array3d, Array4d, ListXd
from ..types import FloatsXd, Ints1d, Ints2d, Ints3d, Ints4d, IntsXd, _Floats
from ..types import DeviceTypes, Generator, Padded, Batchable, SizedGenerator
from ..util import get_array_module, is_xp_array, to_numpy

ArrayT = TypeVar("ArrayT", bound=ArrayXd)
FloatsT = TypeVar("FloatsT", bound=_Floats)
FloatsType = TypeVar("FloatsType", bound=FloatsXd)
SQRT2PI = math.sqrt(2.0 / math.pi)
INV_SQRT2 = 1.0 / math.sqrt(2.0)
INV_SQRT_2PI = 1.0 / math.sqrt(2.0 * math.pi)

class Ops:
    name: str = "base"
    xp: Xp = numpy

    def __init__(
        self, device_type: DeviceTypes = "cpu", device_id: int = -1, **kwargs
    ) -> None:
        self.device_type = device_type
        self.device_id = device_id

    def to_numpy(self, data, *, byte_order=None):  # pragma: no cover
        if isinstance(data, numpy.ndarray):
            if byte_order:
                dtype = data.dtype.newbyteorder(byte_order)
                data = numpy.asarray(data, dtype=dtype)
            return data
        else:
            raise ValueError("Cannot convert non-numpy from base Ops class")

    def minibatch(
        self,
        size: Union[int, Generator],
        sequence: Batchable,
        *,
        shuffle: bool = False,
        buffer: int = 1,
    ) -> SizedGenerator:
        """Iterate slices from a sequence, optionally shuffled. Slices
        may be either views or copies of the underlying data.

        The `size` argument may be either an integer, or a sequence of integers.
        If a sequence, a new size is drawn before every output.

        If shuffle is True, shuffled batches are produced by first generating
        an index array, shuffling it, and then using it to slice into the
        sequence.

        An internal queue of `buffer` items is accumulated before being each
        output. Buffering is useful for some devices, to allow the
        network to run asynchronously without blocking on every batch.
        """
        if not hasattr(sequence, "__len__"):
            err = f"Can't minibatch data. Expected sequence, got {type(sequence)}"
            raise ValueError(err)
        sizes = self._get_batch_sizes(
            len(sequence), itertools.repeat(size) if isinstance(size, int) else size
        )
        indices = numpy.arange(len(sequence))

        # This is a bit convoluted, but it's a time where convenience makes
        # trickery worthwhile: instead of being an actual generator, we
        # return our SizedGenerator object, which provides a __len__.
        def _iter_items():
            if shuffle:
                numpy.random.shuffle(indices)
            queue = []
            i = 0
            for size in sizes:
                size = int(size)
                queue.append(self._get_batch(sequence, indices[i : i + size]))
                if len(queue) >= buffer:
                    yield from queue
                    queue = []
                i += size
            yield from queue

        return SizedGenerator(_iter_items, len(sizes))

    def multibatch(
        self,
        size: Union[int, Generator],
        sequence: Batchable,
        *others: Batchable,
        shuffle: bool = False,
        buffer: int = 1,
    ) -> SizedGenerator:
        """Minibatch one or more sequences of data, and yield
        lists with one batch per sequence. See ops.minibatch.
        """
        # You'd think we could just do this by calling into minibatch and zip...
        # But the shuffling makes it really hard.
        sequences = (sequence,) + tuple(others)
        if not all(hasattr(seq, "__len__") for seq in sequences):
            values = ", ".join([f"{type(seq)}" for seq in sequences])
            err = f"Can't multibatch data. Expected sequences, got {values}"
            raise ValueError(err)
        sizes = self._get_batch_sizes(
            len(sequence), itertools.repeat(size) if isinstance(size, int) else size
        )
        indices = numpy.arange(len(sequence))

        def _iter_items():
            if shuffle:
                numpy.random.shuffle(indices)
            queue = []
            i = 0
            for size in sizes:
                size = int(size)
                idx_batch = indices[i : i + size]
                queue.append([])
                for sequence in sequences:
                    queue[-1].append(self._get_batch(sequence, idx_batch))
                if len(queue) >= buffer:
                    yield from queue
                    queue = []
                i += size
            yield from queue

        return SizedGenerator(_iter_items, len(sizes))

    def _get_batch(self, sequence, indices):
        if isinstance(sequence, list):
            subseq = [sequence[i] for i in indices]
        elif isinstance(sequence, tuple):
            subseq = tuple(sequence[i] for i in indices)
        else:
            subseq = sequence[indices]
        if is_xp_array(subseq):
            subseq = self.as_contig(self.xp.asarray(subseq))
        return subseq

    def _get_batch_sizes(self, length: int, sizes: Iterator[int]):
        output = []
        i = 0
        while i < length:
            output.append(next(sizes))
            i += output[-1]
        return output

    def seq2col(
        self, seq: Floats2d, nW: int, *, lengths: Optional[Ints1d] = None
    ) -> Floats2d:
        """Given an (M, N) sequence of vectors, return an (M, N*(nW*2+1))
        sequence. The new sequence is constructed by concatenating nW preceding
        and succeeding vectors onto each column in the sequence, to extract a
        window of features.
        """
        # This is a test implementation that only supports nW=1 and lengths=None
        assert nW == 1
        assert lengths == None
        B = seq.shape[0]
        I = seq.shape[1]
        cols = self.alloc3f(B, (nW * 2 + 1), I)
        # Copy left contexts. The last words aren't the left-context for anything.
        cols[nW:, :nW] = self.reshape3f(seq[:-nW], -1, nW, I)
        cols[:, nW] = seq
        cols[:-nW, nW + 1 :] = self.reshape3f(seq[nW:], -1, nW, I)
        return self.reshape2f(cols, B, I * (2 * nW + 1))

    def backprop_seq2col(
        self, dY: Floats2d, nW: int, *, lengths: Optional[Ints1d] = None
    ) -> Floats2d:
        """The reverse/backward operation of the `seq2col` function: calculate
        the gradient of the original `(M, N)` sequence, as a function of the
        gradient of the output `(M, N*(nW*2+1))` sequence.
        """
        # This is a test implementation that only supports nW=1 and lengths=None
        assert nW == 1
        assert lengths == None
        nF = nW * 2 + 1
        B = dY.shape[0]
        I = dY.shape[1] // nF
        # Having trouble getting the kernel to work...
        dX = self.alloc2f(B, I)
        dY3d = self.reshape3f(dY, B, nF, I)
        dX[:-nW] += self.reshape2f(dY3d[nW:, :nW], -1, I)
        dX += dY3d[:, nW]
        dX[nW:] += self.reshape2f(dY3d[:-nW, nW + 1 :], -1, I)
        return dX

    def gemm(
        self,
        x: Floats2d,
        y: Floats2d,
        out: Optional[Floats2d] = None,
        trans1: bool = False,
        trans2: bool = False,
    ) -> Floats2d:
        """Perform General Matrix Multiplication (GeMM) and optionally store
        the result in the specified output variable.
        """
        if trans1:
            x = x.T
        if trans2:
            y = y.T
        if out is None:
            return self.xp.dot(x, y)
        else:
            self.xp.dot(x, y, out=out)
            return out

    def tile(self, X: Floats2d, reps: int) -> Floats2d:
        return self.xp.tile(X, reps)

    def affine(self, X: Floats2d, W: Floats2d, b: Floats1d) -> Floats2d:
        """Apply a weights layer and a bias to some inputs, i.e.
        Y = X @ W.T + b
        """
        Y = self.gemm(X, W, trans2=True)
        Y += b
        return Y

    @overload
    def flatten(
        self,
        X: List[Floats2d],
        dtype: Optional[DTypes] = None,
        pad: int = 0,
        ndim_if_empty: int = 2,
     ) -> Floats2d:
        ...

    @overload
    def flatten(
        self,
        X: List[Ints1d],
        dtype: Optional[DTypes] = None,
        pad: int = 0,
        ndim_if_empty: int = 2,
     ) -> Ints1d:
        ...

    @overload
    def flatten(
        self,
        X: List2d,
        dtype: Optional[DTypes] = None,
        pad: int = 0,
        ndim_if_empty: int = 2,
     ) -> Array2d:
        ...

    # further specific typed signatures can be added as necessary

    @overload
    def flatten(
        self,
        X: ListXd,
        dtype: Optional[DTypes] = None,
        pad: int = 0,
        ndim_if_empty: int = 2,
     ) -> ArrayXd:
        ...

    @overload
    def flatten(
        self,
        X: Sequence[ArrayXd],
        dtype: Optional[DTypes] = None,
        pad: int = 0,
        ndim_if_empty: int = 2,
     ) -> ArrayXd:
        ...

    def flatten(
        self,
        X: Sequence[ArrayXd],
        dtype: Optional[DTypes] = None,
        pad: int = 0,
        ndim_if_empty: int = 2,
    ) -> ArrayXd:
        """Flatten a list of arrays into one large array."""
        if X is None or len(X) == 0:
            return self.alloc((0,) * ndim_if_empty, dtype=dtype or "f")
        xp = get_array_module(X[0])
        shape_if_empty = X[0].shape
        X = [x for x in X if x.size != 0]
        if len(X) == 0:
            return self.alloc(shape_if_empty, dtype=dtype or "f")
        if int(pad) >= 1:
            padded = []
            for x in X:
                padded.append(xp.zeros((pad,) + x.shape[1:], dtype=x.dtype))
                padded.append(x)
            padded.append(xp.zeros((pad,) + x.shape[1:], dtype=x.dtype))
            X = padded
        result = xp.concatenate(X)
        if dtype is not None:
            result = xp.asarray(result, dtype=dtype)
        return result

    @overload
    def unflatten(self, X: Floats2d, lengths: Ints1d, pad: int = 0) -> List[Floats2d]:
        ...

    @overload
    def unflatten(self, X: Ints1d, lengths: Ints1d, pad: int = 0) -> List[Ints1d]:
        ...

    @overload
    def unflatten(self, X: Array2d, lengths: Ints1d, pad: int = 0) -> List2d:
        ...

    # further specific typed signatures can be added as necessary

    @overload
    def unflatten(self, X: ArrayXd, lengths: Ints1d, pad: int = 0) -> ListXd:
        ...

    def unflatten(self, X: ArrayXd, lengths: Ints1d, pad: int = 0) -> ListXd:
        """The reverse/backward operation of the `flatten` function: unflatten
        a large array into a list of arrays according to the given lengths.
        """
        # cupy.split requires lengths to be in CPU memory.
        lengths = to_numpy(lengths)

        if pad > 0:
            lengths = numpy.where(lengths > 0, lengths + pad, 0)  # type: ignore
        unflat = self.xp.split(X, numpy.cumsum(lengths))[:-1]  # type: ignore
        if pad > 0:
            unflat = [a[pad:] for a in unflat]

        assert len(unflat) == len(lengths)

        return unflat

    @overload
    def pad(self, seqs: List[Ints2d], round_to=1) -> Ints3d:
        ...

    @overload  # noqa: F811
    def pad(self, seqs: List[Floats2d], round_to=1) -> Floats3d:
        ...

    def pad(  # noqa: F811
        self, seqs: Union[List[Ints2d], List[Floats2d]], round_to=1
    ) -> Array3d:
        """Perform padding on a list of arrays so that they each have the same
        length, by taking the maximum dimension across each axis. This only
        works on non-empty sequences with the same `ndim` and `dtype`.
        """
        # TODO: This should be generalized to handle different ranks
        if not seqs:
            raise ValueError("Cannot pad empty sequence")
        if len(set(seq.ndim for seq in seqs)) != 1:
            raise ValueError("Cannot pad sequences with different ndims")
        if len(set(seq.dtype for seq in seqs)) != 1:
            raise ValueError("Cannot pad sequences with different dtypes")
        if len(set(seq.shape[1:] for seq in seqs)) != 1:
            raise ValueError("Cannot pad sequences that differ on other dimensions")
        # Find the maximum dimension along each axis. That's what we'll pad to.
        length = max(len(seq) for seq in seqs)
        # Round the length to nearest bucket -- helps on GPU, to make similar
        # array sizes.
        length = (length + (round_to - 1)) // round_to * round_to
        final_shape = (len(seqs), length) + seqs[0].shape[1:]
        output: Array3d = self.alloc(final_shape, dtype=seqs[0].dtype)
        for i, arr in enumerate(seqs):
            # It's difficult to convince this that the dtypes will match.
            output[i, : arr.shape[0]] = arr  # type: ignore[assignment, call-overload]
        return output

    def unpad(self, padded: Array3d, lengths: List[int]) -> List2d:
        """The reverse/backward operation of the `pad` function: transform an
        array back into a list of arrays, each with their original length.
        """
        output = []
        for i, length in enumerate(lengths):
            output.append(padded[i, :length])
        return cast(List2d, output)

    def list2padded(self, seqs: List2d) -> Padded:
        """Pack a sequence of 2d arrays into a Padded datatype."""
        if not seqs:
            return Padded(
                self.alloc3f(0, 0, 0), self.alloc1i(0), self.alloc1i(0), self.alloc1i(0)
            )
        elif len(seqs) == 1:
            data = self.reshape3(seqs[0], seqs[0].shape[0], 1, seqs[0].shape[1])
            size_at_t = self.asarray1i([1] * data.shape[0])
            lengths = self.asarray1i([data.shape[0]])
            indices = self.asarray1i([0])
            return Padded(data, size_at_t, lengths, indices)
        lengths_indices = [(len(seq), i) for i, seq in enumerate(seqs)]
        lengths_indices.sort(reverse=True)
        indices_ = [i for length, i in lengths_indices]
        lengths_ = [length for length, i in lengths_indices]
        nS = max([seq.shape[0] for seq in seqs])
        nB = len(seqs)
        nO = seqs[0].shape[1]
        # Reorder the sequences, by length. This looks the same in either
        # direction: you're swapping elements between their original and sorted
        # position.
        seqs = cast(List2d, [seqs[i] for i in indices_])
        arr: Array3d = self.pad(seqs)
        assert arr.shape == (nB, nS, nO), (nB, nS, nO)
        arr = self.as_contig(arr.transpose((1, 0, 2)))
        assert arr.shape == (nS, nB, nO)
        # Build a lookup table so we can find how big the batch is at point t.
        batch_size_at_t_ = [0 for _ in range(nS)]
        current_size = len(lengths_)
        for t in range(nS):
            while current_size and t >= lengths_[current_size - 1]:
                current_size -= 1
            batch_size_at_t_[t] = current_size
        assert sum(lengths_) == sum(batch_size_at_t_)
        return Padded(
            arr,
            self.asarray1i(batch_size_at_t_),
            self.asarray1i(lengths_),
            self.asarray1i(indices_),
        )

    def padded2list(self, padded: Padded) -> List2d:
        """Unpack a Padded datatype to a list of 2-dimensional arrays."""
        data = padded.data
        indices = to_numpy(padded.indices)
        lengths = to_numpy(padded.lengths)
        unpadded: List[Optional[Array2d]] = [None] * len(lengths)
        # Transpose from (length, batch, data) to (batch, length, data)
        data = self.as_contig(data.transpose((1, 0, 2)))
        for i in range(data.shape[0]):
            unpadded[indices[i]] = data[i, : int(lengths[i])]
        return cast(List2d, unpadded)

    def get_dropout_mask(self, shape: Shape, drop: Optional[float]) -> FloatsXd:
        """Create a random mask for applying dropout, with a certain percent of
        the mask (defined by `drop`) will contain zeros. The neurons at those
        positions will be deactivated during training, resulting in a more
        robust network and less overfitting.
        """
        if drop is None or drop <= 0:
            return self.xp.ones(shape, dtype="f")
        elif drop >= 1.0:
            return self.alloc(shape)
        coinflips = self.xp.random.uniform(0.0, 1.0, shape)
        mask = (coinflips >= drop) / (1.0 - drop)
        return cast(FloatsXd, self.asarray(mask, dtype="float32"))

    def alloc1f(
        self,
        d0: int,
        *,
        dtype: Optional[DTypesFloat] = "float32",
        zeros: bool = True,
    ) -> Floats1d:
        return self.alloc((d0,), dtype=dtype, zeros=zeros)

    def alloc2f(
        self,
        d0: int,
        d1: int,
        *,
        dtype: Optional[DTypesFloat] = "float32",
        zeros: bool = True,
    ) -> Floats2d:
        return self.alloc((d0, d1), dtype=dtype, zeros=zeros)

    def alloc3f(
        self,
        d0: int,
        d1: int,
        d2: int,
        *,
        dtype: Optional[DTypesFloat] = "float32",
        zeros: bool = True,
    ) -> Floats3d:
        return self.alloc((d0, d1, d2), dtype=dtype, zeros=zeros)

    def alloc4f(
        self,
        d0: int,
        d1: int,
        d2: int,
        d3: int,
        *,
        dtype: Optional[DTypesFloat] = "float32",
        zeros: bool = True,
    ) -> Floats4d:
        return self.alloc((d0, d1, d2, d3), dtype=dtype, zeros=zeros)

    def alloc_f(
        self,
        shape: Shape,
        *,
        dtype: Optional[DTypesFloat] = "float32",
        zeros: bool = True,
    ) -> FloatsXd:
        return self.alloc(shape, dtype=dtype, zeros=zeros)

    def alloc1i(
        self,
        d0: int,
        *,
        dtype: Optional[DTypesInt] = "int32",
        zeros: bool = True,
    ) -> Ints1d:
        return self.alloc((d0,), dtype=dtype, zeros=zeros)

    def alloc2i(
        self,
        d0: int,
        d1: int,
        *,
        dtype: Optional[DTypesInt] = "int32",
        zeros: bool = True,
    ) -> Ints2d:
        return self.alloc((d0, d1), dtype=dtype, zeros=zeros)

    def alloc3i(
        self,
        d0: int,
        d1: int,
        d2: int,
        *,
        dtype: Optional[DTypesInt] = "int32",
        zeros: bool = True,
    ) -> Ints3d:
        return self.alloc((d0, d1, d2), dtype=dtype, zeros=zeros)

    def alloc4i(
        self,
        d0: int,
        d1: int,
        d2: int,
        d3: int,
        *,
        dtype: Optional[DTypesInt] = "int32",
        zeros: bool = True,
    ) -> Ints4d:
        return self.alloc((d0, d1, d2, d3), dtype=dtype, zeros=zeros)

    def alloc_i(
        self,
        shape: Shape,
        *,
        dtype: Optional[DTypesInt] = "int32",
        zeros: bool = True,
    ) -> IntsXd:
        return self.alloc(shape, dtype=dtype, zeros=zeros)

    def alloc(
        self,
        shape: Shape,
        *,
        dtype: Optional[DTypes] = "float32",
        zeros: bool = True,
    ) -> ArrayT:
        """Allocate an array of a certain shape."""
        if isinstance(shape, int):
            shape = (shape,)

        if zeros:
            return self.xp.zeros(shape, dtype=dtype)
        else:
            return self.xp.empty(shape, dtype=dtype)

    def reshape1(self, array: ArrayXd, d0: int) -> Array1d:
        return cast(Array1d, self.reshape(array, (d0,)))

    def reshape2(self, array: ArrayXd, d0: int, d1: int) -> Array2d:
        return cast(Array2d, self.reshape(array, (d0, d1)))

    def reshape3(self, array: ArrayXd, d0: int, d1: int, d2: int) -> Array3d:
        return cast(Array3d, self.reshape(array, (d0, d1, d2)))

    def reshape4(self, array: ArrayXd, d0: int, d1: int, d2: int, d3: int) -> Array4d:
        return cast(Array4d, self.reshape(array, (d0, d1, d2, d3)))

    def reshape1f(self, array: FloatsXd, d0: int) -> Floats1d:
        return cast(Floats1d, self.reshape(array, (d0,)))

    def reshape2f(self, array: FloatsXd, d0: int, d1: int) -> Floats2d:
        return cast(Floats2d, self.reshape(array, (d0, d1)))

    def reshape3f(self, array: FloatsXd, d0: int, d1: int, d2: int) -> Floats3d:
        return cast(Floats3d, self.reshape(array, (d0, d1, d2)))

    def reshape4f(
        self, array: FloatsXd, d0: int, d1: int, d2: int, d3: int
    ) -> Floats4d:
        return cast(Floats4d, self.reshape(array, (d0, d1, d2, d3)))

    def reshape_f(self, array: FloatsXd, shape: Shape) -> FloatsXd:
        return self.reshape(array, shape)

    def reshape1i(self, array: IntsXd, d0: int) -> Ints1d:
        return cast(Ints1d, self.reshape(array, (d0,)))

    def reshape2i(self, array: IntsXd, d0: int, d1: int) -> Ints2d:
        return cast(Ints2d, self.reshape(array, (d0, d1)))

    def reshape3i(self, array: IntsXd, d0: int, d1: int, d2: int) -> Ints3d:
        return cast(Ints3d, self.reshape(array, (d0, d1, d2)))

    def reshape4i(self, array: IntsXd, d0: int, d1: int, d2: int, d3: int) -> Ints4d:
        return cast(Ints4d, self.reshape(array, (d0, d1, d2, d3)))

    def reshape_i(self, array: IntsXd, shape: Shape) -> IntsXd:
        return self.reshape(array, shape)

    def reshape(self, array: ArrayT, shape: Shape) -> ArrayT:
        """Reshape an array."""
        if isinstance(shape, int):
            shape = (shape,)
        return cast(ArrayT, array.reshape(shape))

    def asarray4f(
        self,
        data: Union[Floats4d, Sequence[int]],
        *,
        dtype: Optional[DTypes] = "float32",
    ) -> Floats4d:
        return cast(Floats4d, self.asarray(data, dtype=dtype))

    def asarray3f(
        self,
        data: Union[Floats3d, Sequence[int]],
        *,
        dtype: Optional[DTypes] = "float32",
    ) -> Floats3d:
        return cast(Floats3d, self.asarray(data, dtype=dtype))

    def asarray2f(
        self,
        data: Union[Floats2d, Sequence[int]],
        *,
        dtype: Optional[DTypes] = "float32",
    ) -> Floats2d:
        return cast(Floats2d, self.asarray(data, dtype=dtype))

    def asarray1f(
        self,
        data: Union[Floats1d, Sequence[int]],
        *,
        dtype: Optional[DTypes] = "float32",
    ) -> Floats1d:
        return cast(Floats1d, self.asarray(data, dtype=dtype))

    def asarray_f(
        self,
        data: Union[FloatsXd, Sequence[float]],
        *,
        dtype: Optional[DTypes] = "float32",
    ) -> FloatsXd:
        return cast(FloatsXd, self.asarray(data, dtype=dtype))

    def asarray1i(
        self, data: Union[Ints1d, Sequence[int]], *, dtype: Optional[DTypes] = "int32"
    ) -> Ints1d:
        return cast(Ints1d, self.asarray(data, dtype=dtype))

    def asarray2i(
        self, data: Union[Ints2d, Sequence[int]], *, dtype: Optional[DTypes] = "int32"
    ) -> Ints2d:
        return cast(Ints2d, self.asarray(data, dtype=dtype))

    def asarray3i(
        self, data: Union[Ints3d, Sequence[int]], *, dtype: Optional[DTypes] = "int32"
    ) -> Ints3d:
        return cast(Ints3d, self.asarray(data, dtype=dtype))

    def asarray4i(
        self, data: Union[Ints4d, Sequence[int]], *, dtype: Optional[DTypes] = "int32"
    ) -> Ints4d:
        return cast(Ints4d, self.asarray(data, dtype=dtype))

    def asarray_i(
        self, data: Union[IntsXd, Sequence[int]], *, dtype: Optional[DTypes] = "int32"
    ) -> IntsXd:
        return cast(IntsXd, self.asarray(data, dtype=dtype))

    def asarray(
        self,
        data: Union[ArrayXd, Sequence[ArrayXd], Sequence[float], Sequence[int]],
        *,
        dtype: Optional[DTypes] = None,
    ) -> ArrayXd:
        """Ensure a given array is of the correct type."""
        if isinstance(data, self.xp.ndarray):
            if dtype is None:
                return data
            elif data.dtype == dtype:
                return data
            else:
                return self.xp.asarray(data, dtype=dtype)
        elif hasattr(data, "numpy"):
            # Handles PyTorch Tensor
            return data.numpy()  # type: ignore[union-attr]
        elif dtype is not None:
            return self.xp.array(data, dtype=dtype)
        else:
            return self.xp.array(data)

    def as_contig(self, data: ArrayT, dtype: Optional[DTypes] = None) -> ArrayT:
        """Allow the backend to make a contiguous copy of an array.
        Implementations of `Ops` do not have to make a copy or make it
        contiguous if that would not improve efficiency for the execution engine.
        """
        if data.flags["C_CONTIGUOUS"] and dtype in (None, data.dtype):
            return data
        kwargs = {"dtype": dtype} if dtype is not None else {}
        return self.xp.ascontiguousarray(data, **kwargs)

    def sigmoid(self, X: FloatsType, *, inplace: bool = False) -> FloatsType:
        # To prevent overflows and help with regularization/numerical stability
        X = self.xp.clip(X, -20.0, 20.0)

        if inplace:
            self.xp.exp(-X, out=X)
            X += 1.0  # type: ignore[assignment]
            X **= -1.0  # type: ignore[assignment]
            return cast(FloatsType, X)
        else:
            return cast(FloatsType, 1.0 / (1.0 + self.xp.exp(-X)))

    def backprop_sigmoid(
        self, dY: FloatsType, Y: FloatsType, *, inplace: bool = False
    ) -> FloatsType:
        if inplace:
            self.dsigmoid(Y, inplace=True)
            Y *= dY  # type: ignore
            return Y
        else:
            return dY * self.dsigmoid(Y, inplace=inplace)  # type: ignore

    def dsigmoid(self, Y: FloatsType, *, inplace: bool = False) -> FloatsType:
        if inplace:
            Y *= 1 - Y
            return Y
        else:
            return Y * (1.0 - Y)

    def dtanh(self, Y: FloatsT, *, inplace: bool = False) -> FloatsT:
        if inplace:
            Y **= 2
            Y *= -1.0
            Y += 1.0
            return Y
        else:
            return 1 - Y**2

    def softmax(
        self,
        x: FloatsT,
        *,
        inplace: bool = False,
        axis: int = -1,
        temperature: float = 1.0,
    ) -> FloatsT:
        if temperature != 1.0:
            x = x / temperature
        maxes = self.xp.max(x, axis=axis, keepdims=True)
        shifted = x - maxes
        new_x = self.xp.exp(shifted)
        new_x /= new_x.sum(axis=axis, keepdims=True)
        return new_x

    def softmax_sequences(
        self, Xs: Floats2d, lengths: Ints1d, *, inplace: bool = False, axis: int = -1
    ) -> Floats2d:
        if Xs.ndim >= 3:
            err = f"Softmax currently only supports 2d. Got: {Xs.ndim}"
            raise NotImplementedError(err)
        # This loses almost no fidelity, and helps the numerical stability.
        Xs = self.xp.clip(Xs, -20.0, 20.0)
        new_x = self.xp.exp(Xs)
        summed = self.backprop_reduce_sum(self.reduce_sum(new_x, lengths), lengths)
        new_x /= summed
        return new_x

    def backprop_softmax(
        self, Y: FloatsT, dY: FloatsT, *, axis: int = -1, temperature: float = 1.0
    ) -> FloatsT:
        if temperature != 1.0:
            dY = dY / temperature

        dX = Y * dY
        dX -= Y * dX.sum(axis=axis, keepdims=True)
        return dX

    def backprop_softmax_sequences(
        self, dY: Floats2d, Y: Floats2d, lengths: Ints1d
    ) -> Floats2d:
        dX = Y * dY
        sum_dX = self.backprop_reduce_sum(self.reduce_sum(dX, lengths), lengths)
        dX -= Y * sum_dX
        return dX

    def lstm_forward_training(
        self,
        params: Floats1d,
        H0: Floats3d,
        C0: Floats3d,
        X: Floats2d,
        size_at_t: Ints1d,
    ) -> Tuple[Floats2d, Tuple]:
        assert H0.shape == C0.shape
        assert H0.shape[1] == C0.shape[1]
        Y, fwd_state = lstm_forward_training(params, H0, C0, X, size_at_t)
        return Y, fwd_state

    def lstm_forward_inference(
        self,
        params: Floats1d,
        H0: Floats3d,
        C0: Floats3d,
        X: Floats2d,
        size_at_t: Ints1d,
    ) -> Floats2d:
        Y, _ = lstm_forward_training(params, H0, C0, X, size_at_t)
        return Y

    def backprop_lstm(
        self, dY: Floats2d, lengths: Ints1d, params: Floats1d, fwd_state: Tuple
    ) -> Tuple[Floats2d, Floats1d]:
        dX, d_params = backprop_lstm(dY, lengths, params, fwd_state)
        return dX, d_params

    def maxout(self, X: Floats3d) -> Tuple[Floats2d, Ints2d]:
        which = X.argmax(axis=-1)
        return X.max(axis=-1), which

    def backprop_maxout(self, dY: Floats2d, which: Ints2d, P: int) -> Floats3d:
        dX = self.alloc3f(dY.shape[0], dY.shape[1], P, dtype=dY.dtype)
        for b in range(dY.shape[0]):
            for o in range(dY.shape[1]):
                dX[b, o, which[b, o]] = dY[b, o]
        return dX

    def relu(self, X: Floats2d, inplace: bool = False) -> Floats2d:
        if not inplace:
            return X * (X > 0)
        else:
            X *= X > 0
            return X

    def backprop_relu(
        self, dY: Floats2d, Y: Floats2d, inplace: bool = False
    ) -> Floats2d:
        if not inplace:
            return dY * (Y > 0)
        dY *= Y > 0
        return dY

    def clipped_linear(
        self,
        X: FloatsType,
        slope: float = 1.0,
        offset: float = 0.0,
        min_val: float = 0.0,
        max_val: float = 1.0,
        inplace: bool = False,
    ) -> FloatsType:
        if inplace:
            X *= slope  # type: ignore[assignment]
            X += offset  # type: ignore[assignment]
            return cast(FloatsType, self.xp.clip(X, min_val, max_val, out=X))
        out = X * slope + offset  # type: ignore[assignment]
        return cast(FloatsType, self.xp.clip(out, min_val, max_val))

    def backprop_clipped_linear(
        self,
        dY: FloatsType,
        X: FloatsType,
        slope: float = 1.0,
        offset: float = 0.0,
        min_val: float = 0.0,
        max_val: float = 1.0,
        inplace: bool = False,
    ) -> FloatsType:
        low = (min_val - offset) / slope
        high = (max_val - offset) / slope
        slope = self.xp.float64(slope).astype(X.dtype)
        zero = self.xp.float64(0.0).astype(X.dtype)
        dX = self.xp.where((low < X) & (X < high), slope, zero)
        if inplace:
            dY *= dX
            return dY
        return dY * dX

    def relu_k(
        self, X: FloatsType, n: float = 6.0, inplace: bool = False
    ) -> FloatsType:
        return self.clipped_linear(X, max_val=n, inplace=inplace)

    def backprop_relu_k(
        self, dY: FloatsType, X: FloatsType, n: float = 6.0, inplace: bool = False
    ) -> FloatsType:
        return self.backprop_clipped_linear(dY, X, max_val=n, inplace=inplace)

    def hard_sigmoid(self, X: FloatsType, inplace: bool = False) -> FloatsType:
        return self.clipped_linear(X, slope=0.2, offset=0.5)

    def backprop_hard_sigmoid(
        self, dY: FloatsType, X: FloatsType, inplace: bool = False
    ) -> FloatsType:
        return self.backprop_clipped_linear(dY, X, slope=0.2, offset=0.5)

    def hard_tanh(self, X: FloatsType, inplace: bool = False) -> FloatsType:
        return self.clipped_linear(X, min_val=-1.0, max_val=1.0)

    def backprop_hard_tanh(
        self, dY: FloatsType, X: FloatsType, inplace: bool = False
    ) -> FloatsType:
        return self.backprop_clipped_linear(dY, X, min_val=-1.0, max_val=1.0)

    def swish(self, X: FloatsType, inplace: bool = False) -> FloatsType:
        if inplace:
            X *= self.sigmoid(X)  # type: ignore[operator, assignment]
            return cast(FloatsType, X)
        out = X * self.sigmoid(X)  # type: ignore[operator]
        return cast(FloatsType, out)

    def backprop_swish(
        self, dY: FloatsType, X: FloatsType, Y: FloatsType, inplace: bool = False
    ) -> FloatsType:
        Y = Y + self.sigmoid(X) * (1 - Y)  # type: ignore[operator]
        if inplace:
            dY *= Y  # type: ignore[operator, assignment]
            return cast(FloatsType, dY)
        out = dY * Y  # type: ignore[operator]
        return cast(FloatsType, out)

    # Following https://www.scitepress.org/Papers/2019/74696/74696.pdf
    def hard_swish(self, X: FloatsType, inplace: bool = False) -> FloatsType:
        if inplace:
            X *= self.hard_sigmoid(X)  # type: ignore[operator, assignment]
            return cast(FloatsType, X)
        out = X * self.hard_sigmoid(X)  # type: ignore[operator]
        return cast(FloatsType, out)

    def backprop_hard_swish(
        self, dY: FloatsType, X: FloatsType, inplace: bool = False
    ) -> FloatsType:
        dX = X * 0.4 + 0.5
        dX[X > 2.5] = 1.0
        dX[X < -2.5] = 0
        if inplace:
            dY *= dX
            return dY
        return dY * dX

    # From https://arxiv.org/pdf/1905.02244v5.pdf
    def hard_swish_mobilenet(self, X: FloatsType, inplace: bool = False) -> FloatsType:
        if inplace:
            X *= self.relu_k(X + 3) / 6
            return X
        return X * (self.relu_k(X + 3) / 6)

    def backprop_hard_swish_mobilenet(
        self, dY: FloatsType, X: FloatsType, inplace: bool = False
    ) -> FloatsType:
        dX = (1 / 6) * (X * 2.0 + 3.0)
        dX[X > 3.0] = 1.0
        dX[X < -3.0] = 0
        if inplace:
            dY *= dX
            return dY
        return dX * dY

    # Code snippet taken from:
    # https://www.johndcook.com/blog/2009/01/19/stand-alone-error-function-erf/
    def erf(self, X: FloatsType) -> FloatsType:
        # save the sign of x
        sign = self.xp.sign(X)
        X = self.xp.abs(X)

        a1 = 0.254829592
        a2 = -0.284496736
        a3 = 1.421413741
        a4 = -1.453152027
        a5 = 1.061405429
        p = 0.3275911

        t = 1.0 / (1.0 + p * X)
        y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * self.xp.exp(
            -X * X
        )
        out = sign * y
        out = out.astype(X.dtype)
        return out

    def sechsq(self, X: FloatsType) -> FloatsType:
        return (1 / self.xp.cosh(X)) ** 2

    def gelu_approx(self, X: FloatsType, inplace: bool = False) -> FloatsType:
        tmp = 1.0 + self.xp.tanh(SQRT2PI * (X + 0.044715 * self.xp.power(X, 3)))
        tmp *= 0.5
        tmp = tmp.astype(X.dtype)
        if inplace:
            X *= tmp
            return X
        else:
            Y = self.xp.array(X)
            Y *= tmp
            return Y

    def backprop_gelu_approx(
        self, dY: FloatsType, X: FloatsType, inplace: bool = False
    ) -> FloatsType:
        dX = self.alloc_f(X.shape)
        Xp3 = self.xp.power(X, 3)
        tmp = 0.5 * self.xp.tanh(0.0356774 * Xp3 + 0.797885 * X)
        tmp += (0.0535161 * Xp3 + 0.398942 * X) * self.sechsq(
            0.0356774 * Xp3 + 0.797885 * X
        )
        tmp += 0.5
        dX += tmp
        if inplace:
            dY *= dX
            return dY
        return dY * dX

    def gelu(self, X: FloatsType, inplace: bool = False) -> FloatsType:
        # GELU(x) = x · Φ(x)
        cdf = gaussian_cdf(self, X)
        if inplace:
            X *= cdf  # type: ignore[operator, assignment]
            return X
        return X * cdf  # type: ignore[operator, return-value]

    def backprop_gelu(
        self, dY: FloatsType, X: FloatsType, inplace: bool = False
    ) -> FloatsType:
        # GELU'(x) = Φ(x) + x · PDF(x)
        dX = gaussian_cdf(self, X) + X * gaussian_pdf(self, X)  # type: ignore[operator]
        if inplace:
            dY *= dX
            return dY
        return dY * dX

    def mish(
        self, X: FloatsType, threshold: float = 20.0, inplace: bool = False
    ) -> FloatsType:
        tmp = X * self.xp.tanh(self.xp.log(1.0 + self.xp.exp(X)))
        Y = self.xp.where(X >= threshold, X, tmp)
        if inplace:
            X[:] = Y
            return X
        else:
            return Y

    def backprop_mish(
        self,
        dY: FloatsType,
        X: Floats2d,
        threshold: float = 20.0,
        inplace: bool = False,
    ) -> FloatsType:
        if dY.shape != X.shape:
            msg = f"arrays have incompatible shapes: {dY.shape} and {X.shape}"
            raise ValueError(msg)

        xp = get_array_module(X)
        indices = X < threshold
        Xsub = X[indices]
        dYsub = dY[indices]
        omega = 4.0 * (Xsub + 1.0)
        omega += 4.0 * xp.exp(2.0 * Xsub)
        omega += xp.exp(3.0 * Xsub)
        omega += xp.exp(Xsub) * ((4.0 * Xsub) + 6.0)
        delta = xp.exp(Xsub) + 1.0
        delta *= delta
        delta += 1.0
        dXsub = dYsub * ((xp.exp(Xsub) * omega) / (delta**2))
        # Gradient when above threshold will ignore softplus.
        if inplace:
            out = dY
        else:
            out = xp.copy(dY)
        out[indices] = dXsub
        return out

    def update_averages(
        self, ema: FloatsT, weights: FloatsT, t: int, max_decay: float = 0.9999
    ) -> None:
        # Internals for optimizer
        decay = (1.0 + t) / (10.0 + t)
        if decay > max_decay:
            decay = max_decay
        ema -= (1 - decay) * (ema - weights)

    def adam(
        self,
        weights: Floats1d,
        gradient: Floats1d,
        mom1: Floats1d,
        mom2: Floats1d,
        beta1: float,
        beta2: float,
        eps: float,
        learn_rate: float,
        mod_rate: float = 1.0,
    ) -> Tuple[Floats1d, Floats1d, Floats1d, Floats1d]:
        # Internals for optimizer
        mom1 *= beta1
        mom2 *= beta2
        mom1 += gradient * (1.0 - beta1)
        mom2 += gradient * gradient * (1.0 - beta2)
        # Here we assume learn rate is calculated by the caller.
        # cdef weight_t a_t = learn_rate * sqrt(1-beta2**hp.t) / (1-beta1**hp.t);
        weights -= learn_rate * (mom1 / (mod_rate * self.xp.sqrt(mom2) + eps))
        return weights, gradient, mom1, mom2

    def clip_gradient(self, gradient: FloatsT, threshold: float) -> FloatsT:
        # Internals for optimizer
        xp = get_array_module(gradient)
        grad_norm = xp.linalg.norm(gradient)
        if grad_norm >= threshold:
            gradient *= threshold / grad_norm
        return gradient

    def logloss(self, y_true: FloatsT, y_pred: FloatsT) -> float:
        # Currently not used
        log_yp = self.xp.log(y_pred + 1e-8)
        loss = (y_true * log_yp) + (1 - y_true) * self.xp.log((1 - y_pred) + 1e-8)
        return -loss

    def reduce_sum(self, X: Floats2d, lengths: Ints1d) -> Floats2d:
        Y = self.alloc2f(lengths.shape[0], X.shape[1], zeros=False)
        start = 0
        for i, length in enumerate(lengths):
            if length < 0:
                raise ValueError(f"all sequence lengths must be >= 0, got {length}")
            elif start + length > X.shape[0]:
                raise IndexError("lengths must sum up to the number of rows")
            elif length:
                Y[i] = X[start : start + length].sum(axis=0)
                start += length
            else:
                Y[i] = 0.0
        return Y

    def reduce_mean(self, X: Floats2d, lengths: Ints1d) -> Floats2d:
        Y = self.alloc2f(lengths.shape[0], X.shape[1], zeros=False)
        start = 0
        for i, length in enumerate(lengths):
            if length < 0:
                raise ValueError(f"all sequence lengths must be >= 0, got {length}")
            elif start + length > X.shape[0]:
                raise IndexError("lengths must sum up to the number of rows")
            elif length:
                Y[i] = X[start : start + length].mean(axis=0)
            else:
                Y[i] = 0.0
            start += length
        return Y

    def reduce_max(self, X: Floats2d, lengths: Ints1d) -> Tuple[Floats2d, Ints2d]:
        Y = self.alloc2f(lengths.shape[0], X.shape[1], dtype=X.dtype, zeros=False)
        which = self.alloc2i(lengths.shape[0], X.shape[1], zeros=False)
        start = 0
        for i, length in enumerate(lengths):
            if length <= 0:
                raise ValueError(f"all sequence lengths must be > 0, got {length}")
            elif start + length > X.shape[0]:
                raise IndexError("lengths must sum up to the number of rows")
            elif length:
                which[i] = X[start : start + length].argmax(axis=0)
                Y[i] = X[start : start + length].max(axis=0)
            start += length
        return Y, which

    def backprop_reduce_sum(self, d_sums: Floats2d, lengths: Ints1d) -> Floats2d:
        dX = self.alloc2f(
            lengths.sum(), d_sums.shape[1], dtype=d_sums.dtype, zeros=False
        )
        start = 0
        for i, length in enumerate(lengths):
            if length < 0:
                raise ValueError(f"all sequence lengths must be >= 0, got {length}")
            dX[start : start + length] = d_sums[i]
            start += length
        return dX

    def backprop_reduce_mean(self, d_means: Floats2d, lengths: Ints1d) -> Floats2d:
        dX = self.alloc2f(
            lengths.sum(), d_means.shape[1], dtype=d_means.dtype, zeros=False
        )
        start = 0
        for i, length in enumerate(lengths):
            if length < 0:
                raise ValueError(f"all sequence lengths must be >= 0, got {length}")
            dX[start : start + length] = d_means[i] / length
            start += length
        return dX

    def backprop_reduce_max(
        self, d_maxes: Floats2d, which: Ints2d, lengths: Ints1d
    ) -> Floats2d:
        dX = self.alloc2f(lengths.sum(), d_maxes.shape[1], dtype=d_maxes.dtype)
        start = 0
        for i, length in enumerate(lengths):
            if length <= 0:
                raise ValueError(f"all sequence lengths must be > 0, got {length}")

            self.xp.put_along_axis(
                dX[start : start + length], which[i].reshape((1, -1)), d_maxes[i], 0
            )
            start += length
        return dX

    def hash(self, ids: Ints1d, seed: int) -> Ints2d:
        """Hash a sequence of 64-bit keys into a table with 4 32-bit keys, using
        murmurhash3.
        """
        from .numpy_ops import NumpyOps

        numpy_ops = NumpyOps()
        return self.asarray2i(
            numpy_ops.hash(numpy_ops.asarray(ids, dtype="uint64"), seed)
        )

    def ngrams(self, n: int, keys: Ints1d) -> Ints1d:
        from .numpy_ops import NumpyOps

        numpy_ops = NumpyOps()
        return self.asarray1i(
            numpy_ops.ngrams(n, numpy_ops.asarray(keys, dtype="uint64"))
        )

    def position_encode(
        self, N: int, D: int, period: int = 10000, out: Optional[Floats2d] = None
    ) -> Floats2d:
        # Currently internals only
        from .numpy_ops import NumpyOps

        numpy_ops = NumpyOps()
        return self.asarray2f(numpy_ops.position_encode(N, D, period, out))

    def scatter_add(
        self, table: FloatsXd, indices: IntsXd, values: FloatsXd
    ) -> FloatsXd:
        return self.xp.add.at(table, indices, values)

    def insert_into(self, shape, Xs):
        """Maybe don't need this? Just a quicky to get Jax working."""
        output = self.alloc(shape, dtype=Xs[0].dtype)
        for i, x in enumerate(Xs):
            output[i, : x.shape[0]] = x
        return output

"""
LSTM Notation (kind of involved, but made it a lot easier to write)

X: Inputs
Y: Outputs (aka hiddens)
C: Cells
G: Gates (Output of non-linearity, i.e. lstm_gates(X @ W.T)
A: Activations (X @ W.T, before non-linearity)

Imagine we have the input:
batch = [
    ["apple", "banana", "cantaloupe", "date", "elderberry"],
    ["aardvark", "bat", "capybara", "dingo", "elephant"]
]

The input variable X will have one vector per word, so X[0, 1] will be banana's
vector, X[0, 1, 0] will be a float, the first element of that vector.

We're computing an output variable Y of shape (nL, nB, nO), so that Y[0, 1] is
the output variable of banana.

A problem with variables for RNNs is keeping the timesteps straight. It's hard
to distinguish the current, previous, and next timesteps. To solve this problem,
we follow the convention that **we are at timestep 3**.

Additionally, the variables for Y and C are offset by one, as the 0th elements
have the initial hiddens and initial cells. So:

    t=3
    Xt3: The input vectors for 'dingo' and 'date', i.e. X[t]
    Yt3: The output vectors for 'dingo' and 'date', i.e. Y[t+1] (Y is offset.)
    Ct2: The cells calculated at 'c...', that are the input for 'd...'
    Ct3: The cells calculated at 'd...', that are the input for 'e...'
    At3: The activations at 'd...'
    Gt3: The gates at 'd...'
"""

def lstm_forward_training(
    params: Floats1d, c_init: Floats3d, h_init: Floats3d, X: Floats2d, lengths: Ints1d
) -> Tuple[Floats2d, Tuple]:
    xp = get_array_module(params)
    depth, dirs, nO = c_init.shape
    N, nI = X.shape
    batch_size = lengths[0]
    # Preallocate these so we can pass them through for loop.
    G = cast(Floats4d, xp.zeros((depth, dirs, X.shape[0], nO * 4), dtype="f"))
    Y = cast(Floats4d, xp.zeros((depth, dirs, X.shape[0], nO), dtype="f"))
    C = cast(Floats4d, xp.zeros((depth, dirs, X.shape[0], nO), dtype="f"))
    Yt2 = cast(Floats2d, xp.zeros((batch_size, nO), dtype="f"))
    Ct2 = cast(Floats2d, xp.zeros((batch_size, nO), dtype="f"))
    # Compute the start and end indices first.
    indices = []
    start = 0
    for batch_size in lengths:
        indices.append((start, start + batch_size))
        start += batch_size
    params_i = 0
    orig_X = X
    for i in range(depth):
        nI = X.shape[1]
        for d in range(dirs):
            # The inits are shaped (depth, dirs, nO). We add the internal dimension
            # to make them set correctly.
            Yt2 = h_init[i, d].reshape((1, nO))  # type: ignore[assignment]
            Ct2 = c_init[i, d].reshape((1, nO))  # type: ignore[assignment]
            layer_params, params_i = _split_weights(params, i, nO, nI, params_i)
            Wx, Wh, bias = _transpose_weights(layer_params)
            G[i, d] += xp.dot(X, Wx.T)
            G[i, d] += bias
            for start, end in indices if d == 0 else reversed(indices):
                # When we iterate left-to-right, t2 might be longer than t3.
                Yt2 = Yt2[: end - start]
                Ct2 = Ct2[: end - start]
                # But in right-to-left, it's the opposite: t3 can be longer.
                Gt3 = G[i, d, start:end]
                Gt3 = Gt3[: Yt2.shape[0]]
                Gt3 += xp.dot(Yt2, Wh.T)
                Gt3_ = cast(Floats3d, Gt3.reshape((-1, nO, 4)))
                hf = sigmoid(Gt3_[:, :, 0])
                hi = sigmoid(Gt3_[:, :, 1])
                ho = sigmoid(Gt3_[:, :, 2])
                hc = xp.tanh(Gt3_[:, :, 3])
                Ct3 = hf * Ct2
                Ct3 += hi * hc
                # Store results
                Gt3 = (
                    xp.hstack((hf, hi, ho, hc))
                    .reshape((-1, 4, nO))
                    .transpose((0, 2, 1))
                    .reshape((-1, nO * 4))
                )
                # Fix the endpoint to account for shorter slices when iterating
                # reversed. Not 100% sure this is right. If there's a bug, look
                # here?
                end = min(end, start + ho.shape[0])
                Y[i, d, start:end] = xp.tanh(Ct3) * ho
                G[i, d, start:end] = Gt3
                C[i, d, start:end] = Ct3
                # Set the t2 variables to the current t3 variables.
                Ct2 = Ct3
                Yt2 = Y[i, d, start:end]
        H = cast(Floats2d, Y[i].transpose((1, 0, 2)).reshape((N, -1)))
        if dirs == 2:
            H = xp.ascontiguousarray(H)
        X = H
    return H, (Y, G, C, orig_X)

def backprop_lstm(dY: Floats2d, lengths: Ints1d, params: Floats1d, fwd_state: Tuple):
    xp = get_array_module(params)

    Y: Floats4d
    G: Floats4d
    C: Floats4d
    X: Floats2d
    Wx: Floats2d
    Wh: Floats2d
    bias: Floats1d
    dWx: Floats2d
    dWh: Floats2d
    d_bias: Floats1d
    Y, G, C, X = fwd_state
    depth, dirs, N, nO = C.shape
    nI = X.shape[1]
    batch_size = lengths[0]
    # We don't need to store all the cells for all the layers.
    dC = cast(Floats2d, xp.zeros((N, nO), dtype=C.dtype))
    dG = cast(Floats2d, xp.zeros((N, nO * 4), dtype=C.dtype))
    d_params = cast(Floats1d, xp.zeros((params.shape[0],), dtype=params.dtype))
    # Collect the params and slices. It makes it a bit easier to get the indexing
    # right, when we're iterating backwards.
    params_i = 0
    all_layer_params: List[List[Tuple[Tuple[Floats2d, Floats2d, Floats1d], int]]] = []
    for i in range(depth):
        all_layer_params.append([])
        n_inputs = nI if i == 0 else (nO * dirs)
        for d in range(dirs):
            layer_params, params_i = _split_weights(params, i, nO, n_inputs, params_i)
            layer_params = _transpose_weights(layer_params)
            all_layer_params[-1].append((layer_params, params_i))
    params_i = 0
    all_layer_grads: List[List[Tuple[Tuple[Floats2d, Floats2d, Floats1d], int]]] = []
    for i in range(depth):
        all_layer_grads.append([])
        n_inputs = nI if i == 0 else (nO * dirs)
        for d in range(dirs):
            layer_grads, params_i = _split_weights(d_params, i, nO, n_inputs, params_i)
            layer_grads = _transpose_weights(layer_grads)
            all_layer_grads[-1].append((layer_grads, params_i))
    # Similarly, we want to compute the indices first
    indices = []
    start = 0
    for batch_size in lengths:
        indices.append((start, start + batch_size))
        start += batch_size

    Xs = [X] + [
        cast(Floats2d, Y[i].transpose((1, 0, 2)).reshape((N, -1)))
        for i in range(depth - 1)
    ]
    dXs = [xp.zeros((X.shape[0], X.shape[1]), dtype=X.dtype) for X in Xs]
    # Okay, now do the actual looping
    for i in reversed(range(depth)):
        dY3d = cast(Floats3d, dY.reshape((N, dirs, nO)).transpose((1, 0, 2)))
        dX = dXs[i]
        X = Xs[i]
        if dirs >= 2:
            dY3d = xp.ascontiguousarray(dY3d)
        for d in range(dirs):
            Wx, Wh, bias = all_layer_params[i][d][0]
            dWx, dWh, d_bias = all_layer_grads[i][d][0]
            if d == 0:
                start_t3, end_t3 = indices[-1]
                layer_indices = indices[:-1]
                layer_indices.reverse()
            else:
                start_t3, end_t3 = indices[0]
                layer_indices = indices[1:]
            for start_t2, end_t2 in layer_indices:
                size = min(end_t2 - start_t2, end_t3 - start_t3)
                dGt3, dCt2 = backprop_lstm_gates(
                    dY3d[d, start_t3 : start_t3 + size],
                    dC[start_t3 : start_t3 + size],
                    G[i, d, start_t3 : start_t3 + size],
                    C[i, d, start_t3 : start_t3 + size],
                    C[i, d, start_t2 : start_t2 + size],
                )
                # Backprop hidden-to-hidden w.r.t. hidden.
                dY3d[d, start_t2 : start_t2 + size] += dGt3 @ Wh
                # Update iteration variables
                dC[start_t2 : start_t2 + size] = dCt2
                start_t3 = start_t2
                end_t3 = end_t2
            # Backprop input-to-hidden w.r.t. weights.
            dWx += dG.T @ X
            # Backprop hidden-to-hidden w.r.t. weights.
            dWh += dG.T @ Y[i, d]
            # Backprop bias
            d_bias += dG.sum(axis=0)
            # Backprop input-to-hidden w.r.t. input
            dX += dG @ Wx
        dY = dX
    assert dX.shape[1] == X.shape[1]
    grad_parts = []
    for layer_grads in all_layer_grads:
        for dir_grads, _ in layer_grads:
            grad_parts.append(_untranspose_unsplit_weights(dir_grads))
    return dX, xp.concatenate(grad_parts)

def _split_weights(params: Floats1d, i: int, nO: int, nI: int, params_i: int):
    Wx_size = 4 * nO * nI
    bx_size = 4 * nO
    Wh_size = 4 * nO * nO
    bh_size = 4 * nO
    Wx = params[params_i : params_i + Wx_size].reshape((4 * nO, nI))
    params_i += Wx_size
    bx = params[params_i : params_i + bx_size].reshape((4 * nO,))
    params_i += bx_size
    Wh = params[params_i : params_i + Wh_size].reshape((4 * nO, nO))
    params_i += Wh_size
    bh = params[params_i : params_i + bh_size].reshape((4 * nO,))
    params_i += bh_size
    return ((Wx, bx), (Wh, bh)), params_i

def _transpose_weights(params):
    # Transpose the parameters so that the gates are the last dimension. This
    # makes it easier to fuse.
    (Wx, bx), (Wh, bh) = params
    xp = get_array_module(Wx)
    Wx = Wx.reshape((4, -1, Wx.shape[-1]))
    Wx = Wx.transpose((1, 0, 2)).reshape((-1, Wx.shape[-1]))
    bx = bx.reshape((4, -1)).transpose((1, 0)).reshape((-1,))
    Wh = Wh.reshape((4, -1, Wh.shape[-1]))
    Wh = Wh.transpose((1, 0, 2)).reshape((-1, Wh.shape[-1]))
    bh = bh.reshape((4, -1)).transpose((1, 0)).reshape((-1,))
    ascontig = xp.ascontiguousarray
    Wx = ascontig(Wx)
    Wh = ascontig(Wh)
    bias = ascontig(bx) + bh
    return Wx, Wh, bias

def _untranspose_unsplit_weights(params):
    Wx, Wh, bias = params
    xp = get_array_module(Wx)
    nO = Wh.shape[1]
    nI = Wx.shape[1]
    Wx = Wx.reshape((-1, 4, nI)).transpose((1, 0, 2)).reshape((-1, nI))
    Wh = Wh.reshape((-1, 4, nO)).transpose((1, 0, 2)).reshape((-1, nO))
    bias = bias.reshape((-1, 4)).transpose((1, 0)).reshape((-1,))
    zeros = xp.zeros(bias.shape, dtype="f")
    return xp.concatenate((Wx.ravel(), bias, Wh.ravel(), zeros))

def backprop_lstm_gates(
    dYt3: Floats2d, dCt3: Floats2d, Gt3: Floats2d, Ct3: Floats2d, Ct2: Floats2d
) -> Tuple[Floats2d, Floats2d]:
    # See above for notation. Step numbering refers to forward_lstm_gates
    xp = get_array_module(dYt3)
    hf, hi, ho, hc = xp.split(Gt3, 4, axis=-1)
    assert hf.shape[0] == hi.shape[0] == ho.shape[0] == hc.shape[0]
    assert hf.shape[0] == dYt3.shape[0] == dCt3.shape[0] == Ct3.shape[0] == Ct2.shape[0]
    tanhCt3 = xp.tanh(Ct3)
    # 3b: Yt3 = tanhCt3 * ho
    d_ho = dYt3 * tanhCt3
    d_tanhCt3 = dYt3 * ho
    # 3a: tanhCt3 = tanh(Ct3)
    dCt3 += d_tanhCt3 * dtanh(tanhCt3)
    # 2b: Ct3 += hi * hc
    d_hi = dCt3 * hc
    d_hc = dCt3 * hi
    # 2a: Ct3 = hf * Ct2
    d_hf = dCt3 * Ct2
    dCt2 = dCt3 * hf
    d_At3_hc = d_hc * dtanh(hc)  # 1d
    d_At3_ho = d_ho * dsigmoid(ho)  # 1c
    d_At3_hi = d_hi * dsigmoid(hi)  # 1b
    d_At3_hf = d_hf * dsigmoid(hf)  # 1a
    dAt3 = xp.concatenate((d_At3_hf, d_At3_hi, d_At3_ho, d_At3_hc), axis=-1)
    return dAt3, dCt2

def sigmoid(X, out=None):
    xp = get_array_module(X)

    # To prevent overflows and help with regularization/numerical stability
    X = xp.clip(X, -20.0, 20.0)
    return 1.0 / (1.0 + xp.exp(-X))

def dsigmoid(Y: ArrayT) -> ArrayT:
    return Y * (1.0 - Y)

def dtanh(Y: ArrayT) -> ArrayT:
    return 1 - Y**2

def gaussian_cdf(ops: Ops, X: FloatsType) -> FloatsType:
    """Gaussian CDF for distribution with mean 0 and stdev 1."""
    return 0.5 * (1.0 + ops.erf(INV_SQRT2 * X))

def gaussian_pdf(ops: Ops, X: FloatsType) -> FloatsType:
    """Gaussian PDF for distribution with mean 0 and stdev 1."""
    return INV_SQRT_2PI * ops.xp.exp(-0.5 * X * X)