Update on "NJT unsqueeze() fixes"

jbschlosser · jbschlosser · commit 5d7f369534eb · 2024-11-26T13:31:56.000-05:00
This PR contains three `unsqueeze()`-related fixes for NJT: 1. Adjusts the output's `_ragged_idx` when `unsqueeze()` inserts a dim before the ragged dim 2. Corrects the unbind reference for `unsqueeze()` after the last input dim. For this case, the dim kwarg canonicalization logic needs to be applied wrt `inp.dim() + 1` to account for `dim=-1` properly 3. Adds ragged dim support to `unsqueeze()`, allowing for e.g. `(B, j1, D) -> (B, 1, j1, D)`. This is okay now after #137125 Note that `unsqueeze()` still doesn't support batch dim operation, and arguably should never support this. [ghstack-poisoned]
diff --git a/torch/testing/_internal/opinfo/definitions/nested.py b/torch/testing/_internal/opinfo/definitions/nested.py
@@ -1385,6 +1385,8 @@ def sample_inputs_where(op_info, device, dtype, requires_grad, **kwargs):
         sample.kwargs["condition"] = sample.input > 0.0
         sample.name = sample.name.replace("(", "(NT, ")
         yield sample
+
+
 # === END OP-SPECIFIC SAMPLE INPUTS FUNCS / REFERENCES ===
 
 
