functorch/dim/op_properties.py

# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import torch
# pointwise operators can go through a faster pathway

tensor_magic_methods = [
    'add',
    ''
]
pointwise_magic_methods_with_reverse = (
    'add', 'sub', 'mul', 'floordiv', 'div', 'truediv', 'mod',
    'pow', 'lshift', 'rshift', 'and', 'or', 'xor'
)
pointwise_magic_methods = (
    *(x for m in pointwise_magic_methods_with_reverse for x in (m, 'r' + m)),
    'eq', 'gt', 'le', 'lt', 'ge', 'gt', 'ne', 'neg', 'pos',
    'abs', 'invert',
    'iadd', 'isub', 'imul', 'ifloordiv', 'idiv',
    'itruediv', 'imod', 'ipow', 'ilshift', 'irshift', 'iand',
    'ior', 'ixor',
    'int', 'long', 'float', 'complex',
)

pointwise_methods = (
    *(f'__{m}__' for m in pointwise_magic_methods),
)

pointwise = (
    *(getattr(torch.Tensor, m) for m in pointwise_methods),
    torch.nn.functional.dropout,
    torch.where,
    torch.Tensor.abs,
    torch.abs,
    torch.Tensor.acos,
    torch.acos,
    torch.Tensor.acosh,
    torch.acosh,
    torch.Tensor.add,
    torch.add,
    torch.Tensor.addcdiv,
    torch.addcdiv,
    torch.Tensor.addcmul,
    torch.addcmul,
    torch.Tensor.addr,
    torch.addr,
    torch.Tensor.angle,
    torch.angle,
    torch.Tensor.asin,
    torch.asin,
    torch.Tensor.asinh,
    torch.asinh,
    torch.Tensor.atan,
    torch.atan,
    torch.Tensor.atan2,
    torch.atan2,
    torch.Tensor.atanh,
    torch.atanh,
    torch.Tensor.bitwise_and,
    torch.bitwise_and,
    torch.Tensor.bitwise_left_shift,
    torch.bitwise_left_shift,
    torch.Tensor.bitwise_not,
    torch.bitwise_not,
    torch.Tensor.bitwise_or,
    torch.bitwise_or,
    torch.Tensor.bitwise_right_shift,
    torch.bitwise_right_shift,
    torch.Tensor.bitwise_xor,
    torch.bitwise_xor,
    torch.Tensor.ceil,
    torch.ceil,
    torch.celu,
    torch.nn.functional.celu,
    torch.Tensor.clamp,
    torch.clamp,
    torch.Tensor.clamp_max,
    torch.clamp_max,
    torch.Tensor.clamp_min,
    torch.clamp_min,
    torch.Tensor.copysign,
    torch.copysign,
    torch.Tensor.cos,
    torch.cos,
    torch.Tensor.cosh,
    torch.cosh,
    torch.Tensor.deg2rad,
    torch.deg2rad,
    torch.Tensor.digamma,
    torch.digamma,
    torch.Tensor.div,
    torch.div,
    torch.dropout,
    torch.nn.functional.dropout,
    torch.nn.functional.elu,
    torch.Tensor.eq,
    torch.eq,
    torch.Tensor.erf,
    torch.erf,
    torch.Tensor.erfc,
    torch.erfc,
    torch.Tensor.erfinv,
    torch.erfinv,
    torch.Tensor.exp,
    torch.exp,
    torch.Tensor.exp2,
    torch.exp2,
    torch.Tensor.expm1,
    torch.expm1,
    torch.feature_dropout,
    torch.Tensor.float_power,
    torch.float_power,
    torch.Tensor.floor,
    torch.floor,
    torch.Tensor.floor_divide,
    torch.floor_divide,
    torch.Tensor.fmod,
    torch.fmod,
    torch.Tensor.frac,
    torch.frac,
    torch.Tensor.frexp,
    torch.frexp,
    torch.Tensor.gcd,
    torch.gcd,
    torch.Tensor.ge,
    torch.ge,
    torch.nn.functional.gelu,
    torch.nn.functional.glu,
    torch.Tensor.gt,
    torch.gt,
    torch.Tensor.hardshrink,
    torch.hardshrink,
    torch.nn.functional.hardshrink,
    torch.nn.functional.hardsigmoid,
    torch.nn.functional.hardswish,
    torch.nn.functional.hardtanh,
    torch.Tensor.heaviside,
    torch.heaviside,
    torch.Tensor.hypot,
    torch.hypot,
    torch.Tensor.i0,
    torch.i0,
    torch.Tensor.igamma,
    torch.igamma,
    torch.Tensor.igammac,
    torch.igammac,
    torch.Tensor.isclose,
    torch.isclose,
    torch.Tensor.isfinite,
    torch.isfinite,
    torch.Tensor.isinf,
    torch.isinf,
    torch.Tensor.isnan,
    torch.isnan,
    torch.Tensor.isneginf,
    torch.isneginf,
    torch.Tensor.isposinf,
    torch.isposinf,
    torch.Tensor.isreal,
    torch.isreal,
    torch.Tensor.kron,
    torch.kron,
    torch.Tensor.lcm,
    torch.lcm,
    torch.Tensor.ldexp,
    torch.ldexp,
    torch.Tensor.le,
    torch.le,
    torch.nn.functional.leaky_relu,
    torch.Tensor.lerp,
    torch.lerp,
    torch.Tensor.lgamma,
    torch.lgamma,
    torch.Tensor.log,
    torch.log,
    torch.Tensor.log10,
    torch.log10,
    torch.Tensor.log1p,
    torch.log1p,
    torch.Tensor.log2,
    torch.log2,
    torch.nn.functional.logsigmoid,
    torch.Tensor.logical_and,
    torch.logical_and,
    torch.Tensor.logical_not,
    torch.logical_not,
    torch.Tensor.logical_or,
    torch.logical_or,
    torch.Tensor.logical_xor,
    torch.logical_xor,
    torch.Tensor.logit,
    torch.logit,
    torch.Tensor.lt,
    torch.lt,
    torch.Tensor.maximum,
    torch.maximum,
    torch.Tensor.minimum,
    torch.minimum,
    torch.nn.functional.mish,
    torch.Tensor.mvlgamma,
    torch.mvlgamma,
    torch.Tensor.nan_to_num,
    torch.nan_to_num,
    torch.Tensor.ne,
    torch.ne,
    torch.Tensor.neg,
    torch.neg,
    torch.Tensor.nextafter,
    torch.nextafter,
    torch.Tensor.outer,
    torch.outer,
    torch.polar,
    torch.Tensor.polygamma,
    torch.polygamma,
    torch.Tensor.positive,
    torch.positive,
    torch.Tensor.pow,
    torch.pow,
    torch.Tensor.prelu,
    torch.prelu,
    torch.nn.functional.prelu,
    torch.Tensor.rad2deg,
    torch.rad2deg,
    torch.Tensor.reciprocal,
    torch.reciprocal,
    torch.Tensor.relu,
    torch.relu,
    torch.nn.functional.relu,
    torch.nn.functional.relu6,
    torch.Tensor.remainder,
    torch.remainder,
    torch.Tensor.round,
    torch.round,
    torch.rrelu,
    torch.nn.functional.rrelu,
    torch.Tensor.rsqrt,
    torch.rsqrt,
    torch.rsub,
    torch.selu,
    torch.nn.functional.selu,
    torch.Tensor.sgn,
    torch.sgn,
    torch.Tensor.sigmoid,
    torch.sigmoid,
    torch.nn.functional.sigmoid,
    torch.Tensor.sign,
    torch.sign,
    torch.Tensor.signbit,
    torch.signbit,
    torch.nn.functional.silu,
    torch.Tensor.sin,
    torch.sin,
    torch.Tensor.sinc,
    torch.sinc,
    torch.Tensor.sinh,
    torch.sinh,
    torch.nn.functional.softplus,
    torch.nn.functional.softshrink,
    torch.Tensor.sqrt,
    torch.sqrt,
    torch.Tensor.square,
    torch.square,
    torch.Tensor.sub,
    torch.sub,
    torch.Tensor.tan,
    torch.tan,
    torch.Tensor.tanh,
    torch.tanh,
    torch.nn.functional.tanh,
    torch.threshold,
    torch.nn.functional.threshold,
    torch.trapz,
    torch.Tensor.true_divide,
    torch.true_divide,
    torch.Tensor.trunc,
    torch.trunc,
    torch.Tensor.xlogy,
    torch.xlogy,
    torch.rand_like,
)