def forward(ctx, inputs, first_idxs, num_inputs):
"""
Args:
ctx: Context object used to calculate gradients.
inputs: FloatTensor of shape (N, max_size, D), representing
the padded tensor, e.g. areas for faces in a batch of meshes.
first_idxs: LongTensor of shape (N,) where N is the number of
elements in the batch and `first_idxs[i] = f`
means that the inputs for batch element i begin at `inputs_packed[f]`.
num_inputs: Number of packed entries (= F)
Returns:
inputs_packed: FloatTensor of shape (F, D) where
`inputs_packed[first_idx[i]:] = inputs[i, :]`.
"""
if not (inputs.dim() == 3):
raise ValueError("input can only be 3-dimensional.")
if not (first_idxs.dim() == 1):
raise ValueError("first_idxs can only be 1-dimensional.")
if not (inputs.dtype == torch.float32):
raise ValueError("input has to be of type torch.float32.")
if not (first_idxs.dtype == torch.int64):
raise ValueError("first_idxs has to be of type torch.int64.")
if not isinstance(num_inputs, int):
raise ValueError("max_size has to be int.")
ctx.save_for_backward(first_idxs)
ctx.max_size = inputs.shape[1]
inputs, first_idxs = inputs.contiguous(), first_idxs.contiguous()
inputs_packed = _C.padded_to_packed(inputs, first_idxs, num_inputs)
return inputs_packed
def backward(ctx, grad_output):
grad_output = grad_output.contiguous()
first_idxs = ctx.saved_tensors[0]
num_inputs = ctx.num_inputs
grad_input = _C.padded_to_packed(grad_output, first_idxs, num_inputs)
return grad_input, None, None
