| from __future__ import annotations | |
| from cupy import _core | |
| from cupy._util import bf16_loop | |
| def create_math_ufunc(math_name, nargs, name, doc, support_complex=True): | |
| assert 1 <= nargs <= 2 | |
| if nargs == 1: | |
| types = ('e->e', *bf16_loop(), 'f->f', 'd->d') | |
| if support_complex: | |
| types += ('F->F', 'D->D') | |
| return _core.create_ufunc( | |
| name, types, 'out0 = %s(in0)' % math_name, doc=doc) | |
| else: | |
| types = ('ee->e', *bf16_loop(2), 'ff->f', 'dd->d') | |
| if support_complex: | |
| types += ('FF->F', 'DD->D') | |
| return _core.create_ufunc( | |
| name, types, 'out0 = %s(in0, in1)' % math_name, doc=doc) | |