58 lines
928 B
Python
58 lines
928 B
Python
|
from torch.masked._ops import (
|
||
|
|
_canonical_dim,
|
||
|
|
_combine_input_and_mask,
|
||
|
|
_generate_docstring,
|
||
|
|
_input_mask,
|
||
|
|
_output_mask,
|
||
|
|
_reduction_identity,
|
||
|
|
_where,
|
||
|
|
amax,
|
||
|
|
amin,
|
||
|
|
argmax,
|
||
|
|
argmin,
|
||
|
|
cumprod,
|
||
|
|
cumsum,
|
||
|
|
log_softmax,
|
||
|
|
logaddexp,
|
||
|
|
logsumexp,
|
||
|
|
mean,
|
||
|
|
median,
|
||
|
|
norm,
|
||
|
|
normalize,
|
||
|
|
prod,
|
||
|
|
softmax,
|
||
|
|
softmin,
|
||
|
|
std,
|
||
|
|
sum,
|
||
|
|
var,
|
||
|
|
)
|
||
|
|
from torch.masked.maskedtensor.core import is_masked_tensor, MaskedTensor
|
||
|
|
from torch.masked.maskedtensor.creation import as_masked_tensor, masked_tensor
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"amax",
|
||
|
|
"amin",
|
||
|
|
"argmax",
|
||
|
|
"argmin",
|
||
|
|
"as_masked_tensor",
|
||
|
|
"cumprod",
|
||
|
|
"cumsum",
|
||
|
|
"is_masked_tensor",
|
||
|
|
"log_softmax",
|
||
|
|
"logaddexp",
|
||
|
|
"logsumexp",
|
||
|
|
"masked_tensor",
|
||
|
|
"MaskedTensor",
|
||
|
|
"mean",
|
||
|
|
"median",
|
||
|
|
"norm",
|
||
|
|
"normalize",
|
||
|
|
"prod",
|
||
|
|
"softmax",
|
||
|
|
"softmin",
|
||
|
|
"std",
|
||
|
|
"sum",
|
||
|
|
"var",
|
||
|
|
]
|