Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,560 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import inspect
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Sequence
|
||||
from functools import lru_cache, partial, wraps
|
||||
from itertools import chain
|
||||
from typing import Optional, TYPE_CHECKING, TypeVar, Union
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.export.decomp_utils import CustomDecompTable
|
||||
|
||||
import torch
|
||||
import torch.library
|
||||
from torch._ops import HigherOrderOperator, OperatorBase, OpOverload, OpOverloadPacket
|
||||
from torch._prims_common import CustomOutParamAnnotation
|
||||
from torch._subclasses.functional_tensor import FunctionalTensor
|
||||
from torch.utils import _pytree as pytree
|
||||
|
||||
|
||||
__all__ = [
|
||||
"decomposition_table",
|
||||
"pre_autograd_decomposition_table",
|
||||
"meta_table",
|
||||
"register_decomposition",
|
||||
"get_decompositions",
|
||||
"core_aten_decompositions",
|
||||
"_should_decompose_because_unsafe_op",
|
||||
]
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
# TODO: relax key type here; torch registrations should be possible to; but
|
||||
# right now this type is accurate
|
||||
global_decomposition_table: dict[str, dict[torch._ops.OperatorBase, Callable]] = (
|
||||
defaultdict(dict)
|
||||
)
|
||||
|
||||
decomposition_table = global_decomposition_table["post_autograd"]
|
||||
pre_autograd_decomposition_table = global_decomposition_table["pre_autograd"]
|
||||
meta_table = global_decomposition_table["meta"]
|
||||
|
||||
|
||||
def _should_decompose_because_unsafe_op(op: torch._ops.OperatorBase) -> bool:
|
||||
"""
|
||||
Returns True if the op must always decompose in export/compile tracing system
|
||||
|
||||
In export, we always decompose certain CIA ops that are tagged with
|
||||
maybe_aliasing_or_mutating because we statically need to know if the op is
|
||||
mutating or not. But these CIA ops could have different behaviour in runtime.
|
||||
|
||||
native_batch_norm is a prim op which has a wrong schema and it needs to be replaced
|
||||
with correct schema. But until then, we will force decompose it via this tag.
|
||||
"""
|
||||
if not isinstance(op, torch._ops.OpOverload):
|
||||
return False
|
||||
if torch.Tag.maybe_aliasing_or_mutating in op.tags:
|
||||
return True
|
||||
return op is torch.ops.aten.native_batch_norm.default
|
||||
|
||||
|
||||
def _add_op_to_registry(registry, op, fn):
|
||||
"""
|
||||
This is an internal API for adding an op to the decomposition table.
|
||||
|
||||
If op is OpOverload, it will be added to the registry directly.
|
||||
If op is OpOverloadPacket, all the valid op_overloads in the packet will be added to the registry.
|
||||
"""
|
||||
overloads: list[torch._ops.OperatorBase] = []
|
||||
if isinstance(op, HigherOrderOperator):
|
||||
# There's no concept of overloads for HigherOrderOperator
|
||||
registry[op] = fn
|
||||
return
|
||||
elif isinstance(op, OpOverload):
|
||||
overloads.append(op)
|
||||
else:
|
||||
if not isinstance(op, OpOverloadPacket):
|
||||
raise AssertionError(f"expected OpOverloadPacket, got {type(op)}")
|
||||
for ol in op.overloads():
|
||||
overloads.append(getattr(op, ol))
|
||||
|
||||
for op_overload in overloads:
|
||||
if op_overload in registry:
|
||||
raise RuntimeError(f"duplicate registrations for {op_overload}")
|
||||
# TorchScript dumps a bunch of extra nonsense overloads
|
||||
# which don't have corresponding dispatcher entries, we need
|
||||
# to filter those out, e.g aten.add.float_int
|
||||
if torch._C._dispatch_has_kernel(op_overload.name()):
|
||||
registry[op_overload] = fn
|
||||
|
||||
|
||||
def _convert_out_params(f):
|
||||
out_annotation = f.__annotations__.get("out")
|
||||
|
||||
# If there are no out params, do not wrap the function.
|
||||
if not out_annotation:
|
||||
return f
|
||||
|
||||
# Hack to detect when out is a Tuple. There seems to be no pretty way of doing this
|
||||
if getattr(out_annotation, "__origin__", None) is tuple:
|
||||
sig = inspect.signature(f)
|
||||
out_names = sig.return_annotation._fields
|
||||
# If out is a tuple, we need to register a function that unpacks all the out
|
||||
# elements as this is what native_functions.yaml expects
|
||||
|
||||
@wraps(f)
|
||||
def _fn(*args, **kwargs):
|
||||
out_kwargs = tuple(kwargs.pop(o, None) for o in out_names)
|
||||
# Either all of the out kwargs are set or none of them
|
||||
is_none = out_kwargs[0] is None
|
||||
if not all((o is None) == is_none for o in out_kwargs):
|
||||
raise AssertionError(
|
||||
f"all out kwargs must be set or none of them, got {out_kwargs}"
|
||||
)
|
||||
return f(*args, **kwargs, out=None if is_none else out_kwargs)
|
||||
|
||||
out_params = [
|
||||
inspect.Parameter(
|
||||
o,
|
||||
kind=inspect.Parameter.KEYWORD_ONLY,
|
||||
default=None,
|
||||
annotation=t,
|
||||
)
|
||||
for o, t in zip(out_names, out_annotation.__args__)
|
||||
]
|
||||
# Drop the out parameter and concatenate the new kwargs in the signature
|
||||
params = chain((v for k, v in sig.parameters.items() if k != "out"), out_params)
|
||||
_fn.__signature__ = inspect.Signature( # type: ignore[attr-defined]
|
||||
parameters=params, # type: ignore[arg-type]
|
||||
return_annotation=sig.return_annotation,
|
||||
)
|
||||
# Drop the out parameter and concatenate the new kwargs in the annotations
|
||||
_fn.__annotations__ = {k: v for k, v in f.__annotations__.items() if k != "out"}
|
||||
for o in out_params:
|
||||
_fn.__annotations__[o.name] = o.annotation
|
||||
|
||||
# Propagate that this function is wrapped by `out_wrapper`
|
||||
_fn._torch_decompositions_out_wrapper = f._torch_decompositions_out_wrapper # type: ignore[attr-defined]
|
||||
|
||||
return _fn
|
||||
|
||||
# Alternatively, there may be a single tensor out parameter with a name
|
||||
# other than "out". This will need special treatment and is indicated by an
|
||||
# annotation, which we will remove here so it is not exposed after wrapping.
|
||||
custom_out_param_name = f.__annotations__.pop(CustomOutParamAnnotation, None)
|
||||
if custom_out_param_name:
|
||||
|
||||
@wraps(f)
|
||||
def _fn(*args, **kwargs):
|
||||
out_kwarg = kwargs.pop(custom_out_param_name, None)
|
||||
return f(*args, **kwargs, out=out_kwarg)
|
||||
|
||||
out_param = inspect.Parameter(
|
||||
custom_out_param_name,
|
||||
kind=inspect.Parameter.KEYWORD_ONLY,
|
||||
default=None,
|
||||
annotation=out_annotation,
|
||||
)
|
||||
|
||||
# Drop the out parameter and concatenate the new kwarg in the signature
|
||||
sig = inspect.signature(f)
|
||||
params = chain(
|
||||
(v for k, v in sig.parameters.items() if k != "out"), (out_param,)
|
||||
)
|
||||
_fn.__signature__ = inspect.Signature( # type: ignore[attr-defined]
|
||||
parameters=params, # type: ignore[arg-type]
|
||||
return_annotation=sig.return_annotation,
|
||||
)
|
||||
|
||||
# Drop the out parameter and concatenate the new kwargs in the annotations
|
||||
_fn.__annotations__ = {k: v for k, v in f.__annotations__.items() if k != "out"}
|
||||
_fn.__annotations__[out_param.name] = out_param.annotation
|
||||
|
||||
return _fn
|
||||
|
||||
return f
|
||||
|
||||
|
||||
def register_decomposition(
|
||||
aten_op, registry=None, *, type="post_autograd", unsafe=False
|
||||
) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
|
||||
"""
|
||||
A decorator to register a function as a decomposition to the Python
|
||||
decomposition table. Use it like this::
|
||||
|
||||
@register_decomposition(torch.ops.aten.clamp_min)
|
||||
def clamp_min(x):
|
||||
return torch.clamp(self, min=min)
|
||||
|
||||
If you are writing a new decomposition, consider contributing it
|
||||
directly to PyTorch in torch._decomp.decompositions.
|
||||
|
||||
This API is experimental; we are almost certainly going to extend
|
||||
the API when we make decompositions eligible for use in transforms (e.g.,
|
||||
autograd) and not just backend tracing, where we then need to know if a
|
||||
decomposition can be used to simulate a transform.
|
||||
|
||||
By default, we also will register it to the Meta key of dispatcher,
|
||||
and replace the c++ Meta implementation if there is already one.
|
||||
|
||||
unsafe kwarg is for reuse of this function for registering non-function
|
||||
things
|
||||
"""
|
||||
|
||||
if type not in {"post_autograd", "pre_autograd", "meta"}:
|
||||
raise AssertionError(
|
||||
f"type must be one of post_autograd, pre_autograd, or meta, got {type}"
|
||||
)
|
||||
|
||||
def decomposition_decorator(fn: Callable[_P, _T]) -> Callable[_P, _T]:
|
||||
orig_fn = fn
|
||||
if not unsafe:
|
||||
fn = _convert_out_params(fn)
|
||||
|
||||
nonlocal registry
|
||||
if registry is None:
|
||||
registry = global_decomposition_table[type]
|
||||
|
||||
def register(op):
|
||||
_add_op_to_registry(registry, op, fn)
|
||||
|
||||
# To handle allowing multiple aten_ops at once
|
||||
pytree.tree_map_(register, aten_op)
|
||||
return orig_fn
|
||||
|
||||
return decomposition_decorator
|
||||
|
||||
|
||||
def get_decompositions(
|
||||
aten_ops: Sequence[torch._ops.OperatorBase | OpOverloadPacket],
|
||||
type: str = "post_autograd",
|
||||
) -> dict[torch._ops.OperatorBase, Callable]:
|
||||
"""
|
||||
Retrieve a dictionary of decompositions corresponding to the list of
|
||||
operator overloads and overload packets passed as input. Overload
|
||||
packets will include all decomposed overloads in the packet. If there is
|
||||
no decomposition for a requested operator, it is silently ignored.
|
||||
|
||||
This API is experimental; we are almost certainly going to give an alternate,
|
||||
more recommended formulation, where a user provides the set of operators
|
||||
they know how to implement, and we provide decompositions for everything
|
||||
not in this set.
|
||||
"""
|
||||
if type not in {"post_autograd", "pre_autograd", "meta"}:
|
||||
raise AssertionError(
|
||||
f"type must be one of post_autograd, pre_autograd, or meta, got {type}"
|
||||
)
|
||||
|
||||
registry = global_decomposition_table[type]
|
||||
packets_to_overloads = defaultdict(list)
|
||||
|
||||
for opo in registry:
|
||||
if isinstance(opo, (OpOverload, OpOverloadPacket)):
|
||||
packets_to_overloads[opo.overloadpacket].append(opo)
|
||||
decompositions: dict[torch._ops.OperatorBase, Callable] = {}
|
||||
for op in aten_ops:
|
||||
if isinstance(op, OpOverloadPacket) and op in packets_to_overloads:
|
||||
for op_overload in packets_to_overloads[op]:
|
||||
decompositions[op_overload] = registry[op_overload]
|
||||
elif isinstance(op, (torch._ops.OperatorBase)) and op in registry:
|
||||
decompositions[op] = registry[op]
|
||||
return decompositions
|
||||
|
||||
|
||||
def remove_decompositions(
|
||||
decompositions: dict[torch._ops.OperatorBase, Callable],
|
||||
aten_ops: Sequence[OpOverload | OpOverloadPacket],
|
||||
) -> None:
|
||||
"""
|
||||
Given a dictionary of decompositions obtained from get_decompositions(), removes
|
||||
operators associated with a list of operator overloads and overload packets passed
|
||||
as input. If the decomposition dictionary does not contain a decomposition that is
|
||||
specified to be removed, it is silently ignored.
|
||||
"""
|
||||
for op in aten_ops:
|
||||
if isinstance(op, OpOverloadPacket):
|
||||
for overload_name in op.overloads():
|
||||
opo = getattr(op, overload_name)
|
||||
decompositions.pop(opo, None)
|
||||
elif isinstance(op, OpOverload):
|
||||
decompositions.pop(op, None)
|
||||
|
||||
|
||||
# populate the table
|
||||
import torch._decomp.decompositions
|
||||
import torch._refs
|
||||
|
||||
|
||||
def core_aten_decompositions() -> "CustomDecompTable":
|
||||
from torch.export.exported_program import default_decompositions
|
||||
|
||||
return default_decompositions()
|
||||
|
||||
|
||||
# See NOTE [Core ATen Ops]
|
||||
#
|
||||
# list was copied from torch/_inductor/decomposition.py
|
||||
# excluding decompositions that results in prim ops
|
||||
# Resulting opset of decomposition is core aten ops
|
||||
def _core_aten_decompositions_post_autograd() -> dict[
|
||||
torch._ops.OperatorBase, Callable
|
||||
]:
|
||||
aten = torch.ops.aten
|
||||
return get_decompositions(
|
||||
[
|
||||
aten.addcdiv,
|
||||
aten.addcdiv_,
|
||||
aten.addcmul,
|
||||
aten.addcmul_,
|
||||
aten.addr,
|
||||
aten.affine_grid_generator,
|
||||
aten.alias_copy,
|
||||
aten.all,
|
||||
aten.aminmax,
|
||||
aten.arange.default,
|
||||
aten.arange.start,
|
||||
aten.avg_pool2d_backward,
|
||||
aten.baddbmm,
|
||||
aten.binary_cross_entropy,
|
||||
aten.binary_cross_entropy_backward,
|
||||
aten.binary_cross_entropy_with_logits,
|
||||
aten.block_diag,
|
||||
aten.bernoulli.p,
|
||||
aten.bernoulli.default,
|
||||
aten.celu,
|
||||
aten.celu_,
|
||||
aten.channel_shuffle,
|
||||
aten.clamp_max,
|
||||
aten.clamp_min,
|
||||
aten.col2im,
|
||||
aten.count_nonzero,
|
||||
aten.linalg_cross,
|
||||
aten.cudnn_batch_norm,
|
||||
aten.cudnn_batch_norm_backward,
|
||||
aten.miopen_batch_norm_backward,
|
||||
aten.deg2rad,
|
||||
aten.deg2rad_,
|
||||
aten.detach,
|
||||
aten.diag_embed,
|
||||
aten.diagonal_backward,
|
||||
aten.diagonal_copy,
|
||||
aten.dot,
|
||||
aten.vdot,
|
||||
aten.elu_,
|
||||
aten.elu_backward,
|
||||
aten._embedding_bag,
|
||||
aten.embedding_dense_backward,
|
||||
aten.empty_like,
|
||||
aten._euclidean_dist.default,
|
||||
aten.expand_as,
|
||||
aten.expand_copy,
|
||||
aten.eye,
|
||||
aten.fill,
|
||||
aten.fill_,
|
||||
aten.floor_divide,
|
||||
aten.frac,
|
||||
aten.frac_,
|
||||
aten._fused_moving_avg_obs_fq_helper,
|
||||
aten.gelu_,
|
||||
aten.gelu_backward,
|
||||
aten.glu,
|
||||
aten.glu_backward,
|
||||
aten.hardshrink,
|
||||
aten.hardsigmoid,
|
||||
aten.hardsigmoid_,
|
||||
aten.hardsigmoid_backward,
|
||||
aten.hardswish,
|
||||
aten.hardswish_,
|
||||
aten.hardswish_backward,
|
||||
aten.hardtanh_,
|
||||
aten.hardtanh_backward,
|
||||
aten.hann_window,
|
||||
aten.heaviside,
|
||||
aten.heaviside_,
|
||||
aten.huber_loss,
|
||||
aten.huber_loss_backward,
|
||||
aten.im2col,
|
||||
aten.index_add.out,
|
||||
aten.index_add.default,
|
||||
aten.index_add_,
|
||||
aten.index_copy.out,
|
||||
aten.index_copy.default,
|
||||
aten.index_copy_,
|
||||
aten.index_fill.int_Scalar,
|
||||
aten.index_fill.int_Tensor,
|
||||
aten.index_fill.int_Scalar_out,
|
||||
aten.index_fill.int_Tensor_out,
|
||||
aten.index_fill_,
|
||||
aten.isin,
|
||||
aten.isneginf,
|
||||
aten.isposinf,
|
||||
aten.l1_loss,
|
||||
aten._lazy_clone,
|
||||
aten._test_parallel_materialize,
|
||||
aten.leaky_relu_,
|
||||
aten.leaky_relu_backward,
|
||||
aten.lerp,
|
||||
aten.lerp_,
|
||||
aten.linspace,
|
||||
aten.logaddexp,
|
||||
aten.logaddexp2,
|
||||
aten.logit,
|
||||
aten.logit_,
|
||||
aten.logit_backward,
|
||||
aten.log_sigmoid_backward,
|
||||
aten.log_sigmoid_forward,
|
||||
aten._log_softmax_backward_data,
|
||||
aten.logspace,
|
||||
aten.logsumexp.default,
|
||||
aten.masked_fill,
|
||||
aten.masked_fill_,
|
||||
aten.max_unpool2d,
|
||||
aten.max_unpool3d,
|
||||
aten.mish,
|
||||
aten.mish_,
|
||||
aten.mish_backward,
|
||||
aten.mse_loss,
|
||||
aten.mse_loss_backward,
|
||||
aten.multi_margin_loss,
|
||||
aten.multilabel_margin_loss_forward,
|
||||
aten.mv,
|
||||
aten.mvlgamma,
|
||||
aten.mvlgamma_,
|
||||
aten.nansum,
|
||||
aten.nan_to_num,
|
||||
aten.nan_to_num_,
|
||||
aten.narrow,
|
||||
aten.native_batch_norm_backward,
|
||||
aten.native_dropout_backward,
|
||||
aten.native_group_norm_backward,
|
||||
aten.native_layer_norm_backward,
|
||||
aten._fused_rms_norm,
|
||||
aten._fused_rms_norm_backward,
|
||||
aten.new_empty,
|
||||
aten.new_full,
|
||||
aten.new_ones,
|
||||
aten.new_zeros,
|
||||
aten.nll_loss2d_forward,
|
||||
aten.nll_loss2d_backward,
|
||||
aten.nll_loss_backward,
|
||||
aten.nll_loss_forward,
|
||||
aten.norm.ScalarOpt_dtype,
|
||||
aten.norm.Scalar,
|
||||
aten.norm.ScalarOpt_dim_dtype,
|
||||
aten.norm.ScalarOpt_dim,
|
||||
aten.norm.dtype_out,
|
||||
aten.norm.out,
|
||||
aten.norm.names_dtype_out,
|
||||
aten.norm.names_out,
|
||||
aten.norm.ScalarOpt_dtype_out,
|
||||
aten.norm.Scalar_out,
|
||||
aten.ones,
|
||||
aten.ones_like,
|
||||
aten.pixel_shuffle,
|
||||
aten.pixel_unshuffle,
|
||||
aten._prelu_kernel,
|
||||
aten._prelu_kernel_backward,
|
||||
aten._reshape_alias,
|
||||
aten.rad2deg,
|
||||
aten.rad2deg_,
|
||||
aten.reflection_pad1d,
|
||||
aten.reflection_pad1d_backward,
|
||||
aten.reflection_pad2d,
|
||||
aten.reflection_pad2d_backward,
|
||||
aten.reflection_pad3d,
|
||||
aten.reflection_pad3d_backward,
|
||||
aten.replication_pad1d,
|
||||
aten.replication_pad2d,
|
||||
aten.replication_pad3d,
|
||||
aten.renorm,
|
||||
aten.renorm_,
|
||||
aten.replication_pad2d,
|
||||
aten.resize_as,
|
||||
aten.roll,
|
||||
aten.rot90,
|
||||
aten.rrelu_with_noise,
|
||||
aten.rrelu_with_noise_,
|
||||
aten.rsub,
|
||||
aten._safe_softmax,
|
||||
aten._scaled_dot_product_flash_attention_for_cpu.default,
|
||||
aten.select_backward,
|
||||
aten.select_scatter,
|
||||
aten.sgn,
|
||||
aten.sgn_,
|
||||
aten.sigmoid_backward,
|
||||
aten.silu,
|
||||
aten.silu_,
|
||||
aten.silu_backward.grad_input,
|
||||
aten.silu_backward,
|
||||
aten.sinc,
|
||||
aten.sinc_,
|
||||
aten.slice_backward,
|
||||
aten.smooth_l1_loss,
|
||||
aten.smooth_l1_loss_backward,
|
||||
aten.soft_margin_loss,
|
||||
aten.soft_margin_loss_backward,
|
||||
aten._softmax_backward_data,
|
||||
aten.softplus,
|
||||
aten.softplus_backward,
|
||||
aten.softshrink,
|
||||
aten.special_entr,
|
||||
aten.special_log_ndtr,
|
||||
aten.special_xlog1py,
|
||||
aten.split.Tensor,
|
||||
aten.split_with_sizes_copy,
|
||||
aten.squeeze_copy,
|
||||
aten.squeeze.default,
|
||||
aten.squeeze.dim,
|
||||
aten.std.correction,
|
||||
aten.std.out,
|
||||
aten.std.correction_out,
|
||||
aten.std.names_out,
|
||||
aten.std.correction_names_out,
|
||||
aten.std_mean.correction,
|
||||
aten.std_mean.correction_out,
|
||||
aten.stack,
|
||||
aten.sum.default,
|
||||
aten.sum.out,
|
||||
aten.t,
|
||||
aten.t_copy,
|
||||
aten.take,
|
||||
aten.tanh_backward,
|
||||
aten.threshold,
|
||||
aten.threshold_,
|
||||
aten.threshold_backward,
|
||||
aten.trace,
|
||||
aten.transpose.int,
|
||||
aten.transpose_copy,
|
||||
aten.tril,
|
||||
aten.tril_,
|
||||
aten.triu,
|
||||
aten.triu_,
|
||||
aten.unbind,
|
||||
aten.unfold_backward,
|
||||
aten.unfold_copy,
|
||||
aten._unsafe_index,
|
||||
aten._unsafe_index_put,
|
||||
aten._unsafe_masked_index,
|
||||
aten._unsafe_masked_index_put_accumulate,
|
||||
aten.unsafe_split.Tensor,
|
||||
aten.unsafe_split_with_sizes,
|
||||
aten.unsqueeze_copy,
|
||||
aten._unsafe_view,
|
||||
aten.upsample_linear1d,
|
||||
aten.upsample_bilinear2d.out,
|
||||
aten.upsample_trilinear3d.out,
|
||||
aten.upsample_nearest2d_backward,
|
||||
aten.view_as_complex,
|
||||
aten.xlogy,
|
||||
aten.xlogy_,
|
||||
aten.zero,
|
||||
aten.zero_,
|
||||
aten.zeros,
|
||||
aten.zeros_like,
|
||||
aten._chunk_cat,
|
||||
aten._weight_norm_interface,
|
||||
]
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,344 @@
|
||||
# mypy: allow-untyped-decorators
|
||||
# mypy: allow-untyped-defs
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch._decomp
|
||||
from torch import Tensor
|
||||
from torch._prims_common.wrappers import _maybe_remove_out_wrapper
|
||||
|
||||
|
||||
decomposition_table = torch._decomp.decomposition_table
|
||||
decomposition_table_for_jvp: dict[torch._ops.OperatorBase, Callable] = {}
|
||||
register_decomposition = torch._decomp.register_decomposition
|
||||
aten = torch.ops.aten
|
||||
|
||||
# NOTE: [forward-mode AD decompositions mechanism]
|
||||
#
|
||||
# The mechanism is in VariableType,
|
||||
# IF any inputs have forward grad
|
||||
# AND there is no forward AD formula implemented
|
||||
# AND the functions are actually differentiable
|
||||
# run the decomposition
|
||||
# See run_jit_decomposition_with_args_for_jvp
|
||||
# We currently use python decompositions that we torchscript.
|
||||
#
|
||||
# Note that we would be building the backward graph at the decomposed level
|
||||
# too, but that is OK, because we would've errored out otherwise anyway.
|
||||
#
|
||||
# TODO: The mechanism we are using to register decompositions doesn't
|
||||
# seem to be exclusively used for jvp. So open question here is whether
|
||||
# torch/csrc/jit/runtime/decomposition_registry.cpp is being used for other things.
|
||||
# If that is the case, we may go down the decomposition path unexpectedly
|
||||
# (and possibly produce an unintelligible error) vs erroring out earlier and
|
||||
# printing that the forward AD formula is not implemented.
|
||||
#
|
||||
# The solution to this may be to have an explicitly white list control when
|
||||
# to enable the decomposition.
|
||||
|
||||
|
||||
def maybe_register_decomposition(op):
|
||||
def decorator(f):
|
||||
try:
|
||||
return register_decomposition(op)(f)
|
||||
except Exception:
|
||||
return f
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# Functions where we need a special decomposition for jvp but there's another version that
|
||||
# should be used more generally (ex. for jvp we need to recompute the mean and variance for
|
||||
# the backwards of a normalization function. Without jvp, it should use the saved value)
|
||||
decomposition_table_for_jvp = {}
|
||||
|
||||
|
||||
def register_decomposition_for_jvp(fn):
|
||||
return register_decomposition(fn, registry=decomposition_table_for_jvp)
|
||||
|
||||
|
||||
def _register_jit_decomposition_for_jvp(decomp, use_python=False):
|
||||
if decomp in decomposition_table_for_jvp:
|
||||
decomposition_table_used = decomposition_table_for_jvp
|
||||
elif decomp in decomposition_table:
|
||||
decomposition_table_used = decomposition_table
|
||||
else:
|
||||
raise RuntimeError(f"could not find decomposition for {decomp}")
|
||||
decomp_fn = decomposition_table_used[decomp]
|
||||
|
||||
# `out_wrapper` extends a decompositions signature with
|
||||
# an `out` parameter. However jit will use the unwrapped function's
|
||||
# signature instead so we need to unwrap here to prevent an error
|
||||
decomp_fn = _maybe_remove_out_wrapper(decomp_fn)
|
||||
|
||||
if use_python:
|
||||
decomp_fn = torch.jit.ignore(decomp_fn)
|
||||
sig = inspect.signature(decomp_fn)
|
||||
|
||||
# Create a string wrapping the function from the signature
|
||||
# example output:
|
||||
# def wrapped_decomp(x: torch.Tensor, y: int, z: int):
|
||||
# return decomp_fn(x, y, z)
|
||||
# Thanks copilot!
|
||||
def get_function_def(sig):
|
||||
param_def = [f"{param_str}" for param_str in sig.parameters.values()]
|
||||
param_use = [f"{param_str}" for param_str in sig.parameters]
|
||||
|
||||
return f"def wrapped_decomp({', '.join(param_def)}):\n return decomp_fn({', '.join(param_use)})\n"
|
||||
|
||||
f_str = get_function_def(sig)
|
||||
graph = torch.jit.CompilationUnit(f_str).wrapped_decomp.graph
|
||||
else:
|
||||
graph = torch.jit.script(decomp_fn).graph
|
||||
torch.jit._register_decomposition(decomp, graph)
|
||||
|
||||
|
||||
# The only decompositions here are temporary or hacks for the purposes of jvp
|
||||
|
||||
|
||||
# TODO: do these also belong here?
|
||||
@maybe_register_decomposition(aten.trace.default)
|
||||
def trace(self: Tensor) -> Tensor:
|
||||
return torch.sum(torch.diag(self))
|
||||
|
||||
|
||||
@maybe_register_decomposition(aten.log_sigmoid_forward.default)
|
||||
def log_sigmoid_forward(self: Tensor) -> tuple[Tensor, Tensor]:
|
||||
min = torch.minimum(self.new_zeros(()), self)
|
||||
z = torch.exp(-torch.abs(self))
|
||||
if self.is_cuda or self.is_xpu:
|
||||
buffer = self.new_zeros((0,))
|
||||
else:
|
||||
buffer = z
|
||||
return min - torch.log1p(z), buffer
|
||||
|
||||
|
||||
def recompute_mean_var(
|
||||
input: Tensor, rstd: Tensor, inner_dim_indices: list[int], keepdim: bool
|
||||
):
|
||||
# for most norm decompositions, it will be the same as the core version except for here.
|
||||
# We recompute the mean and variance so that they track gradients through input
|
||||
|
||||
mean = torch.mean(input, dim=inner_dim_indices, keepdim=keepdim)
|
||||
var = torch.var(input, dim=inner_dim_indices, unbiased=False, keepdim=keepdim)
|
||||
eps = torch.pow(1 / rstd, 2) - var # this makes me so sad inside
|
||||
eps = eps.detach()
|
||||
rstd = 1 / torch.sqrt(var + eps)
|
||||
return mean, rstd
|
||||
|
||||
|
||||
@register_decomposition_for_jvp(aten.native_layer_norm_backward)
|
||||
def native_layer_norm_backward(
|
||||
grad_out: Tensor,
|
||||
input: Tensor,
|
||||
normalized_shape: list[int],
|
||||
mean: Tensor,
|
||||
rstd: Tensor,
|
||||
weight: Tensor | None,
|
||||
bias: Tensor | None,
|
||||
output_mask: list[bool],
|
||||
) -> tuple[Tensor | None, Tensor | None, Tensor | None]:
|
||||
input_shape = input.shape
|
||||
input_ndim = input.dim()
|
||||
|
||||
axis = input_ndim - len(normalized_shape)
|
||||
inner_dims = input_shape[axis:]
|
||||
outer_dims = input_shape[:axis]
|
||||
inner_dim_indices = list(range(axis, input_ndim))
|
||||
outer_dim_indices = list(range(axis))
|
||||
|
||||
N = 1
|
||||
for i in inner_dims:
|
||||
N *= i
|
||||
M = 1
|
||||
for i in outer_dims:
|
||||
M *= i
|
||||
if M <= 0 or N <= 0:
|
||||
return (
|
||||
input.new_zeros(input_shape),
|
||||
input.new_zeros(input_shape[axis:]),
|
||||
input.new_zeros(input_shape[axis:]),
|
||||
)
|
||||
|
||||
mean_, rstd_ = recompute_mean_var(input, rstd, inner_dim_indices, keepdim=True)
|
||||
|
||||
x_hat = (input - mean_) * rstd_
|
||||
if weight is not None:
|
||||
grad_x_hat = grad_out * weight
|
||||
else:
|
||||
grad_x_hat = grad_out
|
||||
a = grad_x_hat * N
|
||||
b = torch.sum(grad_x_hat, inner_dim_indices, True)
|
||||
c1 = torch.mul(grad_x_hat, x_hat)
|
||||
c2 = torch.sum(c1, inner_dim_indices, True)
|
||||
c3 = torch.mul(x_hat, c2)
|
||||
inner = a - b - c3
|
||||
|
||||
if output_mask[0]:
|
||||
d_input: Tensor | None = (rstd_ / N) * inner
|
||||
else:
|
||||
d_input = torch.zeros_like(input) # should be None but doesn't work with vjp
|
||||
|
||||
if output_mask[1] and weight is not None:
|
||||
if len(outer_dim_indices) > 0:
|
||||
d_weight: Tensor | None = torch.sum(
|
||||
grad_out * x_hat, outer_dim_indices, False
|
||||
)
|
||||
else:
|
||||
d_weight = grad_out * x_hat
|
||||
elif weight is not None:
|
||||
d_weight = torch.zeros_like(weight) # should be None but doesn't work with vjp
|
||||
else:
|
||||
d_weight = torch.zeros(()) # should be None but doesn't work with vjp
|
||||
|
||||
if output_mask[2] and bias is not None:
|
||||
if len(outer_dim_indices) > 0:
|
||||
d_bias: Tensor | None = torch.sum(grad_out, outer_dim_indices, False)
|
||||
else:
|
||||
d_bias = grad_out.clone()
|
||||
elif bias is not None:
|
||||
d_bias = torch.zeros_like(bias) # should be None but doesn't work with vjp
|
||||
else:
|
||||
d_bias = torch.zeros(()) # should be None but doesn't work with vjp
|
||||
|
||||
return (d_input, d_weight, d_bias)
|
||||
|
||||
|
||||
def prod(x: list[int]):
|
||||
r = 1
|
||||
for i in x:
|
||||
r *= i
|
||||
return r
|
||||
|
||||
|
||||
@register_decomposition_for_jvp(aten.native_batch_norm_backward)
|
||||
def native_batch_norm_backward(
|
||||
grad_out: Tensor,
|
||||
input: Tensor,
|
||||
weight: Tensor | None,
|
||||
running_mean: Tensor | None,
|
||||
running_var: Tensor | None,
|
||||
save_mean: Tensor | None,
|
||||
save_invstd: Tensor | None,
|
||||
train: bool,
|
||||
eps: float,
|
||||
output_mask: list[bool],
|
||||
) -> tuple[Tensor, Tensor | None, Tensor | None]:
|
||||
input_shape = input.shape
|
||||
input_rank = input.dim()
|
||||
if input_rank < 2:
|
||||
raise AssertionError(f"rank of the input must be at least 2, got {input_rank}")
|
||||
|
||||
axis = 1
|
||||
num_features = prod(input_shape) / input_shape[axis] # type: ignore[arg-type]
|
||||
mean = save_mean
|
||||
invstd = save_invstd
|
||||
if train:
|
||||
if save_mean is None or save_invstd is None:
|
||||
raise AssertionError(
|
||||
"when train=True, save_mean and save_invstd are required"
|
||||
)
|
||||
|
||||
reduciton_dims = [0] + list(range(2, input.dim()))
|
||||
if invstd is None:
|
||||
raise AssertionError("invstd must not be None for typing")
|
||||
mean, invstd = recompute_mean_var(input, invstd, reduciton_dims, keepdim=False)
|
||||
else:
|
||||
if running_mean is None or running_var is None:
|
||||
raise AssertionError(
|
||||
"running_mean and running_var must not be None when train=False"
|
||||
)
|
||||
mean = running_mean
|
||||
invstd = torch.rsqrt(running_var + eps)
|
||||
|
||||
if invstd is None or mean is None:
|
||||
raise AssertionError(
|
||||
f"invstd and mean must not be None, got invstd={invstd}, mean={mean}"
|
||||
)
|
||||
|
||||
broadcast_mask = [1] * input_rank
|
||||
broadcast_mask[axis] = input_shape[axis]
|
||||
|
||||
reduction_axes: list[int] = []
|
||||
for i in range(input_rank):
|
||||
if i != axis:
|
||||
reduction_axes.append(i)
|
||||
|
||||
mean = torch.reshape(mean, broadcast_mask)
|
||||
norm = 1.0 / num_features
|
||||
grad_output_sum = torch.sum(grad_out, reduction_axes)
|
||||
dot_p = torch.sum(grad_out * (input - mean), reduction_axes)
|
||||
|
||||
grad_mean = torch.reshape(grad_output_sum * norm, broadcast_mask)
|
||||
proj_scale = torch.reshape(torch.mul(dot_p * norm, invstd * invstd), broadcast_mask)
|
||||
|
||||
if weight is None:
|
||||
grad_scale = torch.reshape(invstd, broadcast_mask) * 1.0
|
||||
else:
|
||||
grad_scale = torch.reshape(invstd * weight, broadcast_mask)
|
||||
|
||||
if train:
|
||||
proj = (input - mean) * proj_scale
|
||||
grad_input = ((grad_out - proj) - grad_mean) * grad_scale
|
||||
else:
|
||||
grad_input = grad_out * grad_scale
|
||||
|
||||
if output_mask[1]:
|
||||
grad_weight = dot_p * invstd
|
||||
elif weight is not None:
|
||||
grad_weight = torch.zeros_like(
|
||||
weight
|
||||
) # should be None but doesn't work with vjp
|
||||
else:
|
||||
grad_weight = torch.zeros(()) # should be None but doesn't work with vjp
|
||||
|
||||
if output_mask[2]:
|
||||
grad_bias = grad_output_sum
|
||||
else:
|
||||
grad_bias = torch.zeros_like(
|
||||
grad_output_sum
|
||||
) # should be None but doesn't work with vjp
|
||||
|
||||
return (grad_input, grad_weight, grad_bias)
|
||||
|
||||
|
||||
@register_decomposition_for_jvp(aten.batch_norm_backward)
|
||||
def batch_norm_backward(
|
||||
grad_out: Tensor,
|
||||
input: Tensor,
|
||||
weight: Tensor,
|
||||
running_mean: Tensor | None,
|
||||
running_var: Tensor | None,
|
||||
save_mean: Tensor | None,
|
||||
save_var: Tensor | None,
|
||||
update: bool,
|
||||
eps: float,
|
||||
output_mask: list[bool],
|
||||
reserve: Tensor,
|
||||
) -> tuple[Tensor, Tensor | None, Tensor | None]:
|
||||
return native_batch_norm_backward(
|
||||
grad_out,
|
||||
input,
|
||||
weight,
|
||||
running_mean,
|
||||
running_var,
|
||||
save_mean,
|
||||
save_var,
|
||||
update,
|
||||
eps,
|
||||
output_mask,
|
||||
)
|
||||
|
||||
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.trace.default, use_python=True)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.nll_loss_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.nll_loss2d_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten._log_softmax_backward_data.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten._softmax_backward_data.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.log_sigmoid_forward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.native_layer_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.native_batch_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.cudnn_batch_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.batch_norm_backward.default)
|
||||
_register_jit_decomposition_for_jvp(torch.ops.aten.miopen_batch_norm_backward.default)
|
||||
@@ -0,0 +1,272 @@
|
||||
# mypy: allow-untyped-decorators
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch._decomp as decomp
|
||||
from torch._decomp import get_decompositions
|
||||
from torch._ops import OpOverload
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
rng_decompositions: dict[str, dict[OpOverload, Callable]] = defaultdict(dict)
|
||||
|
||||
|
||||
def register_rng_decomposition(aten_op):
|
||||
return decomp.register_decomposition(aten_op, rng_decompositions)
|
||||
|
||||
|
||||
def throw_on_non_cuda(device):
|
||||
raise RuntimeError(
|
||||
f"You are trying to functionalize a {device.type} RNG operator but {device.type} does not "
|
||||
f"use Philox/counter-based RNG. Therefore, functionalizing a {device.type} RNG operator is "
|
||||
"not supported. We are discussing the possibility of a Philox-based RNG implementation for CPU."
|
||||
)
|
||||
|
||||
|
||||
# TODO - We have to register many more distributions here, and also higher level
|
||||
# ops like dropout which have fused implementation and can hide the rand inside.
|
||||
@register_rng_decomposition(aten.rand)
|
||||
def rand(shape, dtype=None, layout=torch.strided, device=None, pin_memory=False):
|
||||
if device and device.type != "cuda":
|
||||
throw_on_non_cuda(device)
|
||||
seed, offset = PhiloxStateTracker.get_state_as_tuple()
|
||||
dtype = dtype or torch.float32
|
||||
out, offset_jump = torch.ops.rngprims.philox_rand(
|
||||
shape, seed, offset, None, device, dtype
|
||||
)
|
||||
PhiloxStateTracker.advance_offset(offset_jump)
|
||||
return out
|
||||
|
||||
|
||||
@register_rng_decomposition(aten.rand_like)
|
||||
def rand_like(
|
||||
x: torch.Tensor,
|
||||
dtype=None,
|
||||
layout=None,
|
||||
device=None,
|
||||
pin_memory=False,
|
||||
memory_format=torch.preserve_format,
|
||||
):
|
||||
device = device or x.device
|
||||
if device.type != "cuda":
|
||||
throw_on_non_cuda(device)
|
||||
dtype = dtype or x.dtype
|
||||
seed, offset = PhiloxStateTracker.get_state_as_tuple()
|
||||
out, offset_jump = torch.ops.rngprims.philox_rand(
|
||||
x.shape, seed, offset, None, device, dtype
|
||||
)
|
||||
PhiloxStateTracker.advance_offset(offset_jump)
|
||||
return out
|
||||
|
||||
|
||||
class PhiloxState:
|
||||
"""
|
||||
Represents a PhiloxRngState - (seed, offset) where offset = base_offset +
|
||||
relative_offset. seed and base_offset basically point to the rng state just
|
||||
before tracing starts. relative offset tracks the totally consumed offset at
|
||||
trace time.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.seed = torch.tensor(())
|
||||
self.base_offset = torch.tensor(())
|
||||
self.relative_offset = 0
|
||||
self.offset_advanced_at_least_once = False
|
||||
|
||||
def validate_state(self):
|
||||
if self.seed.numel() == 0 or self.base_offset.numel() == 0:
|
||||
raise AssertionError(
|
||||
f"seed and base_offset must not be empty, got "
|
||||
f"seed.numel()={self.seed.numel()}, base_offset.numel()={self.base_offset.numel()}"
|
||||
)
|
||||
|
||||
def advance_offset(self, consumed_offset):
|
||||
self.offset_advanced_at_least_once = True
|
||||
self.relative_offset = self.relative_offset + consumed_offset
|
||||
|
||||
def set_state(self, seed, base_offset, relative_offset=0):
|
||||
self.seed = seed
|
||||
self.base_offset = base_offset
|
||||
self.relative_offset = relative_offset
|
||||
|
||||
def get_state_as_tuple(self):
|
||||
self.validate_state()
|
||||
return (self.seed, self.base_offset + self.relative_offset)
|
||||
|
||||
def get_state_as_tensor(self):
|
||||
# Only needed because we override get_rng_state.
|
||||
self.validate_state()
|
||||
return torch.stack([self.seed, self.base_offset + self.relative_offset])
|
||||
|
||||
def set_state_from_tensor(self, state):
|
||||
# Only needed because we override set_rng_state.
|
||||
self.seed, self.base_offset = torch.unbind(state)
|
||||
self.relative_offset = 0
|
||||
|
||||
|
||||
class PhiloxStateTracker:
|
||||
"""
|
||||
Singleton class to track the philox rng state during AOT Autograd tracing.
|
||||
For each aot tracing instance, AOT Autograd resets this tracker and keeps
|
||||
track of both forward and backward offsets. At runtime, we only care about
|
||||
the total consumed forward and backward offsets. For dynamic shapes, these
|
||||
offsets are a function of input shapes. Therefore, the AOT generated graphs
|
||||
have additional outputs that compute total consumed forward and backward
|
||||
offsets.
|
||||
"""
|
||||
|
||||
running_state: PhiloxState
|
||||
fwd_state: PhiloxState
|
||||
bwd_state: PhiloxState
|
||||
|
||||
def __enter__(self):
|
||||
PhiloxStateTracker.reset()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_cal, exc_tb):
|
||||
PhiloxStateTracker.reset()
|
||||
|
||||
@classmethod
|
||||
def reset(cls):
|
||||
cls.running_state = PhiloxState()
|
||||
cls.fwd_state = PhiloxState()
|
||||
cls.bwd_state = PhiloxState()
|
||||
|
||||
@classmethod
|
||||
def mark_beginning_of_forward(cls):
|
||||
# Tells the tracker to use fwd_state as the running state
|
||||
cls.running_state = cls.fwd_state
|
||||
|
||||
@classmethod
|
||||
def mark_beginning_of_backward(cls):
|
||||
# Tells the tracker to use bwd_state as the running state
|
||||
cls.running_state = cls.bwd_state
|
||||
|
||||
@classmethod
|
||||
def record_state(cls, seed, offset, mode):
|
||||
# Records the seed and offset tensors. These tensors are used to invoke
|
||||
# the philox_rand functional primitives.
|
||||
if mode == "forward":
|
||||
cls.fwd_state.set_state(seed, offset)
|
||||
cls.mark_beginning_of_forward()
|
||||
else:
|
||||
if mode != "backward":
|
||||
raise AssertionError(f"mode must be 'backward', got {mode}")
|
||||
cls.bwd_state.set_state(seed, offset)
|
||||
|
||||
@classmethod
|
||||
def get_state_as_tensor(cls):
|
||||
# The only reason this exists is because we override get_rng_state and
|
||||
# set_rng_state during tracing. get_rng_state expects a tensor output,
|
||||
# so return (seed, offset) tuple upset other parts of the program like
|
||||
# ctx.saved_tensors.
|
||||
|
||||
# A bad consequence is that if user saves and restores rng state, we
|
||||
# have little bit of ugliness in the generated code, where we first
|
||||
# concat the (seed, offset) to create a tensor for get_rng_state, and
|
||||
# then split it back to get (seed, offset) tuple in set_rng_state.
|
||||
|
||||
# TODO: Investigate if there is be a better way to wrap the tuple in a
|
||||
# false Tensor object, and then desugar it later on.
|
||||
return cls.running_state.get_state_as_tensor()
|
||||
|
||||
@classmethod
|
||||
def get_state_as_tuple(cls):
|
||||
return cls.running_state.get_state_as_tuple()
|
||||
|
||||
@classmethod
|
||||
def set_state_from_tensor(cls, x):
|
||||
# This is only needed because we override set_rng_state. Look at the
|
||||
# comment in get_state_from_tensor method.
|
||||
cls.running_state.set_state_from_tensor(x)
|
||||
|
||||
@classmethod
|
||||
def advance_offset(cls, consumed_offset):
|
||||
cls.running_state.advance_offset(consumed_offset)
|
||||
|
||||
@classmethod
|
||||
def get_current_relative_offset(cls):
|
||||
return cls.running_state.relative_offset
|
||||
|
||||
@staticmethod
|
||||
def multiple_of_4(offset):
|
||||
# torch cuda rng state offset must be a multiple of 4. For inductor, as
|
||||
# we sum up all the numel, the result might not be a multiple of 4. This
|
||||
# method achieves that.
|
||||
return (offset + 3) // 4 * 4
|
||||
|
||||
@classmethod
|
||||
def get_updated_fwd_offset(cls):
|
||||
# Short circuit if no rand ops were observed
|
||||
if not cls.fwd_state.offset_advanced_at_least_once:
|
||||
return cls.fwd_state.base_offset
|
||||
return cls.multiple_of_4(
|
||||
cls.fwd_state.base_offset + cls.fwd_state.relative_offset
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_updated_bwd_offset(cls):
|
||||
# Short circuit if no rand ops were observed
|
||||
if not cls.bwd_state.offset_advanced_at_least_once:
|
||||
return cls.bwd_state.base_offset
|
||||
return cls.multiple_of_4(
|
||||
cls.bwd_state.base_offset + cls.bwd_state.relative_offset
|
||||
)
|
||||
|
||||
|
||||
# Adding more decompositions which eventually use rand_like inside decomps.
|
||||
# Adding these in rng_decompositions ensures the functionalization of rand_like
|
||||
# ops used in these decomps. The list is copied from inductor codebase, which
|
||||
# uses it for similar purpose.
|
||||
#
|
||||
# Caution - These decomps do not have same accuracy as that of eager. However,
|
||||
# we can't just disable them with a config flag like fallback_random, because
|
||||
# for functionalization of rng ops, we have to decompose these ops.
|
||||
extra_random_decomps = get_decompositions(
|
||||
[
|
||||
aten.cauchy,
|
||||
aten.cauchy_,
|
||||
aten.exponential,
|
||||
aten.exponential_,
|
||||
aten.geometric,
|
||||
aten.geometric_,
|
||||
aten.native_dropout,
|
||||
aten.normal,
|
||||
aten.normal_,
|
||||
aten.normal_functional,
|
||||
aten.log_normal,
|
||||
aten.log_normal_,
|
||||
aten.rrelu_with_noise,
|
||||
aten.rrelu_with_noise_,
|
||||
aten.uniform_,
|
||||
]
|
||||
)
|
||||
register_extra_random_decomp = functools.partial(
|
||||
decomp.register_decomposition, registry=extra_random_decomps
|
||||
)
|
||||
|
||||
|
||||
@register_extra_random_decomp([aten.bernoulli_])
|
||||
def bernoulli_(self, p=0.5):
|
||||
if self.device == torch.device("cpu"):
|
||||
return NotImplemented
|
||||
return self.copy_(torch.rand_like(self, dtype=torch.float32) < p)
|
||||
|
||||
|
||||
@register_extra_random_decomp([aten.bernoulli.p])
|
||||
def bernoulli_p(self, p=0.5, *, generator=None):
|
||||
if self.device == torch.device("cpu"):
|
||||
return NotImplemented
|
||||
if generator is not None:
|
||||
raise AssertionError(f"generator must be None, got {generator}")
|
||||
return torch.rand_like(self, dtype=torch.float32) < p
|
||||
|
||||
|
||||
rng_decompositions.update(extra_random_decomps) # type: ignore[arg-type]
|
||||
Reference in New Issue
Block a user