@@ -306,15 +306,28 @@ def map_dim(
306306 return None
307307
308308 groups = self ._valid_groups ()
309- if not self ._is_valid_reduction (normalized_dims , groups .source_axis_to_groups ):
309+ if not self ._is_valid_reduction_or_singleton (
310+ normalized_dims , groups .source_axis_to_groups
311+ ):
310312 return None
311313
312- target_dims = self ._map_dims (
313- normalized_dims ,
314- groups .source_axis_to_groups ,
315- groups .group_to_target_axes ,
314+ source_to_target_axes = self .source_to_target_axes ()
315+ target_dims = sorted (
316+ _dedupe (
317+ target_axis
318+ for source_dim in normalized_dims
319+ for target_axis in source_to_target_axes [source_dim ]
320+ )
316321 )
317- if not target_dims or not self ._is_valid_reduction (
322+ if not target_dims or any (
323+ source_axis not in normalized_dims
324+ for target_axis in target_dims
325+ for source_axis in self .source_axes_for_target_axis (
326+ target_axis , source_to_target_axes
327+ )
328+ ):
329+ return None
330+ if not self ._is_valid_reduction_or_singleton (
318331 target_dims , groups .target_axis_to_groups
319332 ):
320333 return None
@@ -432,6 +445,226 @@ def map_permutation_inverse(
432445 else None
433446 )
434447
448+ def remap_target_shape (self , source_shape : Sequence [_Dim ]) -> list [_Dim ] | None :
449+ if len (source_shape ) != self .source_rank :
450+ return None
451+
452+ source_to_target_axes = self .source_to_target_axes ()
453+ target_to_source_axes = [
454+ self .source_axes_for_target_axis (target_axis , source_to_target_axes )
455+ for target_axis in range (self .target_rank )
456+ ]
457+ target_shape : list [_Dim ] = [1 ] * self .target_rank
458+
459+ for source_axis , target_axes in enumerate (source_to_target_axes ):
460+ updates = self ._target_axis_updates_for_source_axis (
461+ source_shape ,
462+ source_axis ,
463+ target_axes ,
464+ target_to_source_axes ,
465+ )
466+ if updates is None :
467+ return None
468+ for target_axis , target_dim in updates :
469+ target_shape [target_axis ] = target_dim
470+
471+ if not same_numel (source_shape , target_shape ):
472+ return None
473+ if not self ._preserves_source_axis_order (source_shape , source_to_target_axes ):
474+ return None
475+ return target_shape
476+
477+ def _target_axis_updates_for_source_axis (
478+ self ,
479+ source_shape : Sequence [_Dim ],
480+ source_axis : int ,
481+ target_axes : Sequence [int ],
482+ target_to_source_axes : Sequence [Sequence [int ]],
483+ ) -> list [tuple [int , _Dim ]] | None :
484+ if not target_axes :
485+ return []
486+
487+ if len (target_axes ) == 1 :
488+ target_axis = target_axes [0 ]
489+ source_axes = target_to_source_axes [target_axis ]
490+ if source_axis != source_axes [0 ]:
491+ return []
492+ target_dim = numel (source_shape [source_axis ] for source_axis in source_axes )
493+ return [(target_axis , target_dim )]
494+
495+ if any (
496+ len (target_to_source_axes [target_axis ]) > 1 for target_axis in target_axes
497+ ):
498+ return []
499+
500+ target_dims = [self .target_shape [target_axis ] for target_axis in target_axes ]
501+ if _dim_equals (source_shape [source_axis ], self .source_shape [source_axis ]):
502+ return list (zip (target_axes , target_dims ))
503+ if _dim_equals (numel (target_dims ), 1 ):
504+ return [(target_axes [0 ], source_shape [source_axis ])]
505+ if _dim_equals (numel (target_dims ), self .source_shape [source_axis ]):
506+ return list (zip (target_axes , target_dims ))
507+ return None
508+
509+ def remap_unit_slice (
510+ self ,
511+ producer_shape : Sequence [_Dim ],
512+ slice_dim : int ,
513+ start : _Dim ,
514+ end : _Dim ,
515+ step : _Dim = 1 ,
516+ ) -> tuple [list [_Dim ], int , _Dim , _Dim ] | None :
517+ """Move a view before a unit slice.
518+
519+ Returns the new view shape and slice interval for:
520+
521+ view(slice(x, dim, start, end), self.target_shape)
522+ == slice(view(x, new_shape), new_dim, new_start, new_end)
523+
524+ This handles the case where a unit slice produces a singleton source
525+ axis that the view removes, so normal source-to-target dim mapping has
526+ no target axis for the slice dim.
527+
528+ """
529+ if (
530+ len (producer_shape ) != self .source_rank
531+ or not isinstance (slice_dim , int )
532+ or not isinstance (start , (int , torch .SymInt ))
533+ or not isinstance (end , (int , torch .SymInt ))
534+ or not isinstance (step , (int , torch .SymInt ))
535+ ):
536+ return None
537+ if not _dim_equals (step , 1 ) or not _dim_equals (end - start , 1 ):
538+ return None
539+
540+ try :
541+ slice_dim = _normalize_dim (slice_dim , self .source_rank )
542+ except AssertionError :
543+ return None
544+
545+ source_to_target_axes = self .source_to_target_axes ()
546+ if source_to_target_axes [slice_dim ]:
547+ return None
548+
549+ prev_target_axes = [
550+ target_axis
551+ for target_axes in source_to_target_axes [:slice_dim ]
552+ for target_axis in target_axes
553+ ]
554+ next_target_axes = [
555+ target_axis
556+ for target_axes in source_to_target_axes [slice_dim + 1 :]
557+ for target_axis in target_axes
558+ ]
559+ fold_axes = [
560+ target_axes [0 ]
561+ for target_axes in source_to_target_axes [slice_dim + 1 :]
562+ if target_axes
563+ ]
564+ fold_axes = [
565+ target_axis
566+ for target_axis in fold_axes
567+ if all (
568+ prev_target_axis <= target_axis for prev_target_axis in prev_target_axes
569+ )
570+ and all (
571+ target_axis <= next_target_axis for next_target_axis in next_target_axes
572+ )
573+ ]
574+ if not fold_axes :
575+ return None
576+
577+ fold_axis = fold_axes [0 ]
578+ target_shape = list (self .target_shape )
579+ chunk = target_shape [fold_axis ]
580+ target_shape [fold_axis ] = chunk * producer_shape [slice_dim ]
581+ return target_shape , fold_axis , start * chunk , end * chunk
582+
583+ def source_to_target_axes (self ) -> list [list [int ]]:
584+ groups = self ._valid_groups ()
585+ source_to_target_axes = [
586+ self ._map_dims (
587+ [source_axis ],
588+ groups .source_axis_to_groups ,
589+ groups .group_to_target_axes ,
590+ )
591+ for source_axis in range (self .source_rank )
592+ ]
593+
594+ self ._add_singleton_axes (source_to_target_axes )
595+ return source_to_target_axes
596+
597+ def map_source_dims_to_target_axes (
598+ self , source_dims : int | Sequence [int ]
599+ ) -> list [int ] | None :
600+ try :
601+ normalized_dims = _normalize_dims (source_dims , self .source_rank )
602+ except AssertionError :
603+ return None
604+ source_to_target_axes = self .source_to_target_axes ()
605+ return _dedupe (
606+ target_axis
607+ for source_dim in normalized_dims
608+ for target_axis in source_to_target_axes [source_dim ]
609+ )
610+
611+ @staticmethod
612+ def source_axes_for_target_axis (
613+ target_axis : int , source_to_target_axes : Sequence [Sequence [int ]]
614+ ) -> list [int ]:
615+ return [
616+ source_axis
617+ for source_axis , target_axes in enumerate (source_to_target_axes )
618+ if target_axis in target_axes
619+ ]
620+
621+ def _add_singleton_axes (self , source_to_target_axes : list [list [int ]]) -> None :
622+ mapped_source_axes = {
623+ source_axis
624+ for source_axis , target_axes in enumerate (source_to_target_axes )
625+ if target_axes
626+ }
627+ mapped_target_axes = {
628+ target_axis
629+ for target_axes in source_to_target_axes
630+ for target_axis in target_axes
631+ }
632+ source_singletons = [
633+ axis
634+ for axis , dim in enumerate (self .source_shape )
635+ if axis not in mapped_source_axes and _dim_equals (dim , 1 )
636+ ]
637+ target_singletons = [
638+ axis
639+ for axis , dim in enumerate (self .target_shape )
640+ if axis not in mapped_target_axes and _dim_equals (dim , 1 )
641+ ]
642+
643+ if len (source_singletons ) == len (target_singletons ):
644+ pairs = zip (source_singletons , target_singletons )
645+ elif len (source_singletons ) == 1 :
646+ pairs = zip (source_singletons * len (target_singletons ), target_singletons )
647+ elif len (target_singletons ) == 1 :
648+ pairs = zip (source_singletons , target_singletons * len (source_singletons ))
649+ else :
650+ pairs = zip (source_singletons , target_singletons )
651+
652+ for source_axis , target_axis in pairs :
653+ source_to_target_axes [source_axis ].append (target_axis )
654+
655+ @staticmethod
656+ def _preserves_source_axis_order (
657+ source_shape : Sequence [_Dim ],
658+ source_to_target_axes : Sequence [Sequence [int ]],
659+ ) -> bool :
660+ target_axes = [
661+ target_axis
662+ for source_axis , axes in enumerate (source_to_target_axes )
663+ if not _dim_equals (source_shape [source_axis ], 1 )
664+ for target_axis in axes
665+ ]
666+ return target_axes == sorted (target_axes )
667+
435668 @staticmethod
436669 def _map_dims (
437670 source_dims : Iterable [int ],
@@ -553,6 +786,30 @@ def _is_valid_reduction(
553786 group_to_axes [group ].issubset (normalized_dims ) for group in selected_groups
554787 )
555788
789+ @staticmethod
790+ def _is_valid_reduction_or_singleton (
791+ normalized_dims : Iterable [int ],
792+ axis_to_groups : Sequence [Sequence [int ]],
793+ ) -> bool :
794+ """Return whether dims cover complete groups, allowing singleton
795+ axes.
796+ """
797+ normalized_dims = set (normalized_dims )
798+ if not normalized_dims :
799+ return False
800+
801+ group_to_axes : dict [int , set [int ]] = defaultdict (set )
802+ selected_groups : set [int ] = set ()
803+ for axis , groups in enumerate (axis_to_groups ):
804+ for group in groups :
805+ group_to_axes [group ].add (axis )
806+ if axis in normalized_dims :
807+ selected_groups .add (group )
808+
809+ return all (
810+ group_to_axes [group ].issubset (normalized_dims ) for group in selected_groups
811+ )
812+
556813 @classmethod
557814 def _build_groups (
558815 cls , source_shape : Sequence [_Dim ], target_shape : Sequence [_Dim ]
0 commit comments