gen_ufuncs_2.py

from dump_namespace import grab_namespace, get_signature

import numpy as np

namespace = np

dct = grab_namespace(namespace)


# SKIP these (need special handling)
skip = {np.frexp, np.modf,    # non-standard unary ufunc signatures
        np.isnat,
        np.invert,      # bitwise NOT operator
        np.spacing,     # niche, does not have a direct equivalent
}

# np functions where torch names differ
torch_names = {np.radians : "deg2rad",
               np.degrees : "rad2deg",
               np.conjugate : "conj_physical",
               np.fabs : "absolute",       # FIXME: np.fabs raises form complex
               np.rint : "round",
               np.left_shift: "bitwise_left_shift",
               np.right_shift: "bitwise_right_shift",
}


# np functions which do not have a torch equivalent
default_stanza = "torch.{torch_name}(x, out=out)"

stanzas = {np.cbrt : "torch.pow(x, 1/3, out=out)",

           # XXX what on earth is np.positive
           np.positive: "+x",

           # these three do not have an out arg
           np.isinf: "torch.isinf(x)",
           np.isnan: "torch.isnan(x)",
           np.isfinite: "torch.isfinite(x)",
}


# for these np functions, pytorch analog does not have the out= arg
needs_out = {np.isinf, np.isnan, np.isfinite, np.positive}
add_out_stanza = """
    if out is not None:
        out[...] = result
"""


header = """\
# this file is autogenerated via gen_ufuncs.py
# do not edit manually!

import torch

import _util
from _ndarray import asarray_replacer

"""

test_header = header + """\
import numpy as np
import torch

from _unary_ufuncs import *
from testing import assert_allclose
"""


template = """ """

test_template = """

def test_{np_name}():
    assert_allclose(np.{np_name}(0.5),
                    {np_name}(0.5), atol=1e-14, check_dtype=False)

"""


###### UNARY UFUNCS ###################################

_all_list = []
main_text = header
test_text = test_header

_impl_list = []
_ufunc_list = []

for ufunc in dct['ufunc']:
    if ufunc in skip:
        continue

    if ufunc.nin == 1:
#        print(get_signature(ufunc))

        torch_name = torch_names.get(ufunc)
        if torch_name is None:
            torch_name = ufunc.__name__

#        print(ufunc.__name__, ' -- ', torch_name)   

        _impl_stanza = "{np_name} = deco_unary_ufunc(torch.{torch_name})"
        _impl_stanza = _impl_stanza.format(np_name=ufunc.__name__,
                                     torch_name=torch_name,)
        _impl_list.append(_impl_stanza)        

        continue

        torch_stanza = stanzas.get(ufunc)
        if torch_stanza is None:
            torch_stanza = default_stanza.format(torch_name=torch_name)

        out_stanza= add_out_stanza if ufunc in needs_out else ""

        main_text += template.format(np_name=ufunc.__name__,
                                     torch_stanza=torch_stanza,
                                     out_stanza=out_stanza)
        test_text += test_template.format(np_name=ufunc.__name__)

        _all_list.append(ufunc.__name__)


print("\n".join(_impl_list))
print("\n\n-----\n\n")

'''
main_text += "\n\n__all__ = %s" % _all_list


with open("_unary_ufuncs.py", "w") as f:
    f.write(main_text)

with open("test_unary_ufuncs.py", "w") as f:
    f.write(test_text)
'''

###### BINARY UFUNCS ###################################



test_header = header + """\
import numpy as np
import torch

from _binary_ufuncs import *
from testing import assert_allclose
"""


template = """


"""

test_template = """

def test_{np_name}():
    assert_allclose(np.{np_name}(0.5, 0.6),
                               {np_name}(0.5, 0.6), atol=1e-7, check_dtype=False)

"""



skip = {np.divmod,    # two outputs
}


torch_names = {np.power: "pow",
               np.equal: "eq",
}


_all_list = []
main_text = header
test_text = test_header

_impl_list = []
_ufunc_list = []


for ufunc in dct['ufunc']:

    if ufunc in skip:
        continue

    if ufunc.nin == 2:
 ##       print(get_signature(ufunc))

        torch_name = torch_names.get(ufunc)
        if torch_name is None:
            torch_name = ufunc.__name__

        _impl_stanza = "{np_name} = deco_binary_ufunc(torch.{torch_name})"
        _impl_stanza = _impl_stanza.format(np_name=ufunc.__name__,
                                     torch_name=torch_name,)
        _impl_list.append(_impl_stanza)

        _ufunc_stanza = "{np_name} = deco_ufunc_from_impl(_ufunc_impl.{np_name})"
        _ufunc_stanza = _ufunc_stanza.format(np_name=ufunc.__name__)
        _ufunc_list.append(_ufunc_stanza)


print("\n".join(_impl_list))

print("\n\n")
print("\n".join(_ufunc_list))