Package: partialtorch
PartialTorchPartialTorchis a thin C++ wrapper ofPyTorch's operators to support masked and partial semantics.Main FeaturesMasked PairWe use a custom C++ extension class calledpartialtorch.MaskedPairto storedataandmask(an optionalTensorof the same shape asdata, containing0/1values indicating the availability of the corresponding
element indata).The advantages ofMaskedPairis that it is statically-typed but unpackable likenamedtuple,
and more importantly, it is accepted bytorch.jit.scriptfunctions as argument or return type.
This container is a temporary substitution fortorch.masked.MaskedTensorand may change in the future.This table compares the two in some aspects:torch.masked.MaskedTensorpartialtorch.MaskedPairBackendPythonC++NatureIs a subclass ofTensorwithmaskas an additional attributeIs a container ofdataandmaskSupported layoutsStrided and SparseOnly Strided️Mask typestorch.BoolTensorOptional[torch.BoolTensor](may support other dtypes)Ops CoverageListedhere(with lots of restrictions)All masked ops thattorch.masked.MaskedTensorsupports and moretorch.jit.script-ableYes✔️ (Python ops seem not to be jit compiled but encapsulated)Yes✔️SupportsTensor's methodsYes✔️Only a few[^1]Supports__torch_function__Yes✔️No❌[^1]PerformanceSlow and sometimes buggy (e.g. try calling.backward3 times)Faster, not prone to bugs related toautogradas it is a container[^1]: We blametorch😅More details about the differences will be discussed below.Masked OperatorsMasked operatorsare the same things that can be found intorch.maskedpackage (which is, unfortunately, still in prototype stage).Our semantic differs fromtorch.maskedfor non-unary operators.torch.masked: Requires operands to share identical mask
(check thislink), which is not always the case when we have to deal
with missing data.partialtorch: Allows operands to have different masks, the output mask is the result of abitwise allfunction
of input masks' values.Partial OperatorsSimilar to masked operators,partial operatorsallow non-uniform masks but instead of usingbitwise allto compute output mask, they usebitwise any.
That means output at any position with at least one present operand is NOT considered missing.In details, before fowarding to the regulartorchnative operators, the masked positions of each operand are filled
with anidentity value.
The identity value is defined as the initial value that has the propertyop(op_identity, value) = value.
For example, the identity value of element-wise addition is0.All partial operators have a prefixpartial_prepended to their name (e.g.partialtorch.partial_add),
while masked operators inherit their native ops' names.
Reduction operators are excluded from this rule as they can be considered unary partial, and some of them
are already available intorch.masked.Scaled Partial OperatorsSome partial operators that involves addition/substraction are extended to haverescaling semantic.
We call themscaled partial operators.
In essence, they rescale the output by the ratio of present operands in the computation of the output.
The idea is similar totorch.dropoutrescaling by $\frac{1}{1-p}$,
or more precisely the wayPartial Convolutionworks.Programatically, all scaled partial operators share the same signature with their non-scaled counterparts,
and are dispatched to when adding a keyword-only argumentscaled = True:pout=partialtorch.partial_add(pa,pb,scaled=True)Torch Ops CoverageWe found out that the workload is behemoth for a group of one person, and involves manually reimplementing all
native functors under theat::_opsnamespace (guess how many there are).
Therefore, we try to cover as many primitive operators as possible, as well as a few other operators relevant to our
work.
The full list of all registered signatures can be found in thisfile.If you want any operator to be added, please contact me.
But if they fall into one of the following categories, the porting may take long or will not happen:Ops that do not have a meaningful masked semantic (e.g.torch.det).Ops that cannot be implemented easily by calling native ops and requires writing custom kernels (e.g.torch.mode).Ops that accept output as an input a.k.a.outops (e.g.aten::mul.out(self: Tensor, other: Tensor, *, out: Tensor(a!)) -> Tensor(a!)).Ops for tensors with unsuported properties (e.g. named tensors, sparse/quantized layouts).Ops with any input/return type that do not havepybind11type conversions predefined bytorch's C++ backend.Also, everyone is welcome to contribute.Requirementstorch>=2.1.0(this version ofPyTorchbrought a number of changes that are not backward compatible)InstallationFrom TestPyPIpartialtorchhas wheels hosted atTestPyPI(it is not likely to reach a stable state anytime soon):pipinstall-ihttps://test.pypi.org/simple/partialtorchThe Linux and Windows wheels are built withCuda 12.1.
If you cannot find a wheel for your Arch/Python/Cuda, or there is any problem with library linking when importing,
proceed toinstructions to build from source.Linux/WindowsMacOSPython version:3.8-3.113.8-3.11PyTorch version:torch==2.1.0torch==2.1.0Cuda version:12.1-GPU CCs:5.0,6.0,6.1,7.0,7.5,8.0,8.6,9.0+PTX-From SourceFor installing from source, you need a C++17 compiler (gcc/msvc) and a Cuda compiler (nvcc) installed.
Then, clone this repo and execute:pipinstall.UsageInitializing aMaskedPairWhileMaskedPairis almost as simple as anamedtuple, there are also a few supporting creation ops:importtorch,partialtorchx=torch.rand(3,3)x_mask=torch.bernoulli(torch.full_like(x,0.5)).bool()# x_mask must have dtype torch.boolpx=partialtorch.masked_pair(x,x_mask)# with 2 inputs data and maskpx=partialtorch.masked_pair(x)# with data only (mask = None)px=partialtorch.masked_pair(x,None)# explicitly define mask = Nonepx=partialtorch.masked_pair(x,True)# explicitly define mask = True (equivalent to None)px=partialtorch.masked_pair((x,x_mask))# from tuple# this new random function conveniently does the work of the above stepspx=partialtorch.rand_mask(x,0.5)Note thatMaskedPairis not a subclass ofTensorlikeMaskedTensor,
so we only support a very limited number of methods.
This is mostly because of the current limitations of C++ backend for custom classes[^1] such as:Unable to overload methods with the same nameUnable to define custom type conversions from Python type (Tensor) or to custom Python type
(to be able to define custom methods such as__str__ofTensordoes for example)Unable to define__torch_function__In the meantime, please considerMaskedPairpurely a fast container and usepartialtorch.op(pair, ...)instead ofpair.op(...)if not available.Note:You cannot indexMaskedPairwithpair[..., 1:-1]as they acts like tuple of 2 elements when indexed.OperatorsAll registered ops can be accessed like any torch's custom C++ operator by callingtorch.ops.partialtorch.[op_name](the same way we call native ATen functiontorch.ops.aten.[op_name]).
Their overloaded versions that acceptTensorare also registered for convenience
(but return type is always converted toMaskedPair).torchpartialtorchimporttorchtorch.manual_seed(1)x=torch.rand(5,5)y=torch.sum(x,0,keepdim=True)importtorchimportpartialtorchtorch.manual_seed(1)x=torch.rand(5,5)px=partialtorch.rand_mask(x,0.5)# standard extension ops callingpout=torch.ops.partialtorch.sum(px,0,keepdim=True)# all exposed ops are also aliased inside partialtorch.opspout=partialtorch.ops.sum(px,0,keepdim=True)Furthermore, we inherit the naming convention of for inplace ops - appending a trailing_character after their
names (e.g.partialtorch.reluandpartialtorch.relu_).
They modify both data and mask of the first operand inplacely.The usage is kept as close to the correspondingTensorops as possible.
Hence, further explaination is redundant.Neural Network LayersCurrently, there are only a number of modules implemented inpartialtorch.nnsubpackage that are masked equivalences
of those intorch.nn.
This is the list of submodules insidepartialtorch.nn.modulesand the layers they provide:partialtorch.nn.modules.activation: All activations
excepttorch.nn.MultiheadAttentionpartialtorch.nn.modules.batchnorm:BatchNormNdpartialtorch.nn.modules.channelshuffle:ChannelShufflepartialtorch.nn.modules.conv:PartialConvNd,PartialConvTransposeNdpartialtorch.nn.modules.dropout:DropoutNd,AlphaDropout,FeatureAlphaDropoutpartialtorch.nn.modules.flatten:Flatten,Unflattenpartialtorch.nn.modules.fold:Fold,Unfoldpartialtorch.nn.modules.instancenorm:InstanceNormNdpartialtorch.nn.modules.normalization:LayerNormpartialtorch.nn.modules.padding:CircularPadNd,ConstantPadNd,ReflectionPadNd,ReplicationPadNd,ZeroPadNdpartialtorch.nn.modules.pixelshuffle:PixelShuffle,PixelUnshufflepartialtorch.nn.modules.pooling:MaxPoolNd,AvgPoolNd,FractionalMaxPoolNd,LpPoolNd,AdaptiveMaxPoolNd,AdaptiveAvgPoolNdpartialtorch.nn.modules.upsampling:Upsample,UpsamplingNearest2d,UpsamplingBilinear2d,PartialUpsample,PartialUpsamplingBilinear2dThe steps for declaring your custom module is identical, except that we now use the classes insidepartialtorch.nnwhich input and outputMaskedPair.
Note that to make them scriptable, you may have to explicitly annotate input and output types.torchpartialtorchimporttorch.nnasnnimporttorch.nn.functionalasFfromtorchimportTensorclassConvBlock(nn.Module):def__init__(self,in_channels,out_channels):self.conv=nn.Conv2d(in_channels,out_channels,kernel_size=(3,3))self.bn=nn.BatchNorm2d(out_channels)self.pool=nn.MaxPool2d(kernel_size=(2,2))defforward(self,x:Tensor)->Tensor:x=self.conv(x)x=F.relu(x)x=self.bn(x)x=self.pool(x)returnximporttorch.nnasnnimportpartialtorch.nnaspartial_nnimportpartialtorch.nn.functionalaspartial_FfrompartialtorchimportMaskedPairclassPartialConvBlock(nn.Module):def__init__(self,in_channels,out_channels):self.conv=partial_nn.PartialConv2d(in_channels,out_channels,kernel_size=(3,3))self.bn=partial_nn.BatchNorm2d(out_channels)self.pool=partial_nn.MaxPool2d(kernel_size=(2,2))defforward(self,x:MaskedPair)->MaskedPair:x=self.conv(x)x=partial_F.relu(x)x=self.bn(x)x=self.pool(x)returnxA few other examples can be found inexamplesfolder.CitationThis code is part of another project of us. Citation will be added in the future.AcknowledgementsPart of the codebase is modified from the following repositories:https://github.com/pytorch/pytorchhttps://github.com/NVIDIA/partialconvLicenseThe code is released under the MIT license. SeeLICENSE.txtfor details.