def unique(tensor: storch.Tensor, event_dim: Optional[int] = 0) -> storch.Tensor:
with storch.ignore_wrapping():
fl_tensor = torch.flatten(tensor, tensor.plate_dims)
uniq, inverse_indexing = torch.unique(
fl_tensor, return_inverse=True, dim=event_dim
)
inverse_indexing = storch.Tensor(
inverse_indexing, [tensor], tensor.plates, "inv_index_" + tensor.name
)
uq_plate = UniquePlate(
"uq_plate_" + tensor.name,
uniq.shape[0],
tensor.multi_dim_plates(),
inverse_indexing,
)
return storch.Tensor(uniq, [tensor], [uq_plate], "unique_" + tensor.name)
def on_unwrap_tensor(self, tensor: storch.Tensor) -> storch.Tensor:
"""
Gets called whenever the given tensor is being unwrapped and unsqueezed for batch use.
This method should not be called on tensors whose variable index is higher than this plates.
:param tensor: The input tensor that is being unwrapped
:return: The tensor that will be unwrapped and unsqueezed in the future. Can be a modification of the input tensor.
"""
if self._in_recursion:
# Required when calling storch.gather in this method. It will call on_unwrap_tensor again.
return tensor
for i, plate in enumerate(tensor.multi_dim_plates()):
if plate.name != self.name:
continue
assert isinstance(plate, AncestralPlate)
if plate.variable_index == self.variable_index:
return tensor
# This is true by the filtering at on_collecting_args
assert plate.variable_index < self.variable_index
parent_plates = []
current_plate = self
# Collect the list of plates from the tensors variable index to this plates variable index
while current_plate.variable_index != plate.variable_index:
parent_plates.append(current_plate)
current_plate = current_plate.parent_plate
assert current_plate == plate
# Go over all parent plates and gather their respective choices.
for parent_plate in reversed(parent_plates):
self._in_recursion = True
expanded_selected_samples = expand_with_ignore_as(
parent_plate.selected_samples, tensor, self.name
)
self._override_equality = False
# Gather what samples of the tensor are chosen by this plate (parent_plate)
tensor = storch.gather(
tensor, parent_plate.name, expanded_selected_samples
)
self._in_recursion = False
self._override_equality = False
break
return tensor