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