Source code for earthkit.utils.array.namespace.torch

# (C) Copyright 2025 ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

from earthkit.utils.array.namespace.unknown import UnknownPatchedNamespace
from earthkit.utils.decorators import thread_safe_cached_property




class PatchedTorchNamespace(UnknownPatchedNamespace):


    def __init__(self):
        super().__init__(None)


    @thread_safe_cached_property
    def xp(self):
        import array_api_compat.torch as torch

        return torch

    @property
    def _earthkit_array_namespace_name(self):
        return "torch"



    def sign(self, x, *args, **kwargs):
        """Reimplement the sign function to handle NaNs.

        The problem is that torch.sign returns 0 for NaNs, but the array API
        standard requires NaNs to be propagated.
        """
        r = self.xp.sign(x, *args, **kwargs)
        r[self.xp.isnan(x)] = self.xp.nan
        return r




    def percentile(self, a, q, axis=None):
        return self.xp.quantile(a, q / 100, dim=axis)




    def quantile(self, a, q, axis=None):
        return self.xp.quantile(a, q, dim=axis)




    def size(self, x):
        """Return the size of an array."""
        return x.numel()




    def shape(self, x):
        """Return the shape of an array."""
        return tuple(x.shape)




    def histogramdd(self, x, *, bins=10):
        return self.xp.histogramdd(x, bins=bins)




    def to_device(self, x, device, **kwargs):
        return x.to(device, **kwargs)




    def rad2deg(self, x):
        return self.xp.rad2deg(x)




    def deg2rad(self, x):
        return self.xp.deg2rad(x)




    def choice(self, a, size, replace=True, generator=None):
        kwargs = {}
        if not isinstance(a, int):
            kwargs["device"] = self.device(a)
        rng = self.xp.Generator(**kwargs) if generator is None else generator
        return self.xp.multinomial(a, size, replacement=replace, generator=rng)