returnn 1.20250901.123052__py3-none-any.whl → 1.20260105.192646__py3-none-any.whl

returnn/frontend/conv.py

@@ -3,7 +3,7 @@ Convolution, transposed convolution, pooling
 """
 
 from __future__ import annotations
-from typing import Optional, Sequence, Tuple, Union
+from typing import Optional, Union, TypeVar, Sequence, Tuple, List
 from returnn.util.basic import next_type_attrib_in_mro_chain
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
@@ -25,6 +25,9 @@ __all__ = [
     "pool2d",
     "pool3d",
     "make_conv_out_spatial_dims",
+    "calc_conv_out_length",
+    "make_transposed_conv_out_spatial_dims",
+    "calc_transposed_conv_out_length",
 ]
 
 
@@ -396,7 +399,11 @@ def transposed_conv(
             )
         if use_mask:
             source = source.copy_masked(0, dims=in_spatial_dims)
-    if padding == "same" and _any_is_non_default(strides, default=1) and _should_use_consistent_same_padding():
+    if (
+        padding == "same"
+        and any(s != 1 for s in (strides or [fs.dimension for fs in filter_size]))
+        and _should_use_consistent_same_padding()
+    ):
         # I don't really know what this should mean here... Investigate this further...
         raise NotImplementedError("consistent same padding not implemented for transposed conv")
     # noinspection PyProtectedMember
@@ -424,6 +431,39 @@ class TransposedConv1d(_TransposedConv):
 
     nd = 1
 
+    def __init__(
+        self,
+        in_dim: Dim,
+        out_dim: Dim,
+        filter_size: Union[int, Dim],
+        *,
+        padding: str,
+        remove_padding: int = 0,
+        output_padding: Optional[int] = None,
+        strides: Optional[int] = None,
+        with_bias: bool = True,
+    ):
+        """
+        :param in_dim:
+        :param out_dim:
+        :param filter_size:
+        :param strides: specifies the upscaling. by default, same as filter_size
+        :param padding: "same" or "valid"
+        :param remove_padding:
+        :param output_padding:
+        :param with_bias: whether to add a bias. enabled by default
+        """
+        super().__init__(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            filter_size=[filter_size],
+            padding=padding,
+            remove_padding=remove_padding,
+            output_padding=output_padding,
+            strides=[strides] if strides is not None else None,
+            with_bias=with_bias,
+        )
+
     __call__ = _ConvOrTransposedConv._call_nd1
 
 
@@ -704,7 +744,7 @@ def make_conv_out_spatial_dims(
     strides: Union[Sequence[int], int] = 1,
     dilation_rate: Union[Sequence[int], int] = 1,
     description_prefix: Optional[str] = None,
-) -> Sequence[Dim]:
+) -> List[Dim]:
     """create out spatial dims from in spatial dims"""
     nd = len(in_spatial_dims)
     if isinstance(filter_size, (int, Dim)):
@@ -715,84 +755,263 @@ def make_conv_out_spatial_dims(
         strides = [strides] * nd
     if isinstance(dilation_rate, int):
         dilation_rate = [dilation_rate] * nd
-    assert nd == len(in_spatial_dims) == len(filter_size) == len(strides) == len(dilation_rate)
     if isinstance(padding, (int, str)):
         padding = [padding] * nd
+    assert nd == len(in_spatial_dims) == len(filter_size) == len(strides) == len(dilation_rate) == len(padding)
     padding = [p.lower() if isinstance(p, str) else p for p in padding]
     out_spatial_dims = []
     for i in range(nd):
-        in_spatial_dim = in_spatial_dims[i]
-        if (filter_size[i] == strides[i] == 1 and padding[i] in ("valid", "same", 0)) or (
-            strides[i] == 1 and padding[i] == "same"
-        ):
-            out_spatial_dims.append(in_spatial_dim)
-        else:
-            out_spatial_dim = _calc_out_dim(
-                in_dim=in_spatial_dim,
+        out_spatial_dims.append(
+            calc_conv_out_length(
+                in_spatial_dims[i],
                 filter_size=filter_size[i],
+                padding=padding[i],
                 stride=strides[i],
                 dilation_rate=dilation_rate[i],
-                padding=padding[i],
+                name=f"{description_prefix}:spatial{i}" if description_prefix else None,
             )
-            assert isinstance(out_spatial_dim, Dim)
-            if description_prefix and out_spatial_dim != in_spatial_dim:
-                out_spatial_dim.name = f"{description_prefix}:spatial{i}"
-            if in_spatial_dim.dyn_size_ext is not None and out_spatial_dim.dyn_size_ext is None:
-                out_spatial_dim.dyn_size_ext = _calc_out_dim(
-                    in_dim=in_spatial_dim.dyn_size_ext,
-                    filter_size=filter_size[i],
-                    stride=strides[i],
-                    dilation_rate=dilation_rate[i],
-                    padding=padding[i],
-                )
-            out_spatial_dims.append(out_spatial_dim)
+        )
     return out_spatial_dims
 
 
-def _calc_out_dim(in_dim, filter_size, stride, padding, dilation_rate=1):
+T = TypeVar("T", int, Dim, Tensor)
+
+
+def calc_conv_out_length(
+    in_length: Union[T, int, Dim, Tensor],
+    *,
+    filter_size: Union[T, int, Dim, Tensor],
+    stride: int,
+    padding: Union[str, int],
+    dilation_rate: int = 1,
+    name: Optional[str] = None,
+) -> T:
     """
     Copied and adapted from TF ConvLayer.calc_out_dim.
 
-    :param T|int|Tensor|torch.Tensor|tensorflow.Tensor|Dim in_dim: dimension in some axis
-    :param int filter_size: e.g. 2, for the corresponding axis
-    :param int stride: e.g. 1, for the corresponding axis
-    :param int dilation_rate: e.g. 1
-    :param str|int padding: "valid" or "same" or int
+    :param T in_length: dimension in some axis
+    :param filter_size: e.g. 2, for the corresponding axis
+    :param stride: e.g. 1, for the corresponding axis
+    :param dilation_rate: e.g. 1
+    :param padding: "valid" or "same" or int
+    :param name:
     :return: the output dimension
-    :rtype: T
     """
+    padding = padding.lower() if isinstance(padding, str) else padding
+    if isinstance(filter_size, int):
+        filter_size_int = filter_size
+    elif isinstance(filter_size, Dim):
+        filter_size_int = filter_size.dimension
+    else:
+        filter_size_int = None
+    filter_size_ = filter_size_int if isinstance(filter_size_int, int) else filter_size
 
-    def ceildiv(a, b):
-        """
-        :param T|int|Tensor|torch.Tensor|tensorflow.Tensor a:
-        :param T|int|Tensor|torch.Tensor|tensorflow.Tensor b:
-        :rtype: T
-        """
-        if isinstance(b, int) and b == 1:
-            return a
-        if isinstance(a, Tensor):
-            return rf.ceil_divide(a, b)
-        return -(-a // b)
+    if (filter_size_int == stride == 1 and padding in ("valid", "same", 0)) or (stride == 1 and padding == "same"):
+        return in_length
 
-    padding = padding.lower() if isinstance(padding, str) else padding
     # See tf.compat.v1.nn.convolution() documentation for more.
     if padding == "same":
-        if isinstance(in_dim, Dim):
-            return in_dim.ceildiv_right(stride)
-        return ceildiv(in_dim, stride)
+        if isinstance(in_length, Dim):
+            out_length = in_length.ceildiv_right(stride)
+        else:
+            out_length = _ceildiv(in_length, stride)
     elif padding == "valid" or isinstance(padding, int):
         if isinstance(padding, int) and padding != 0:
             assert padding > 0
-            in_dim = padding + in_dim + padding
-        if isinstance(in_dim, Dim):
-            filter_left_dilated = (filter_size - 1) * dilation_rate // 2
-            filter_right_dilated = (filter_size - 1) * dilation_rate - filter_left_dilated
-            valid_part = in_dim.sub_left(filter_left_dilated).sub_right(filter_right_dilated)
-            return valid_part.ceildiv_right(stride)
-        return ceildiv(in_dim - (filter_size - 1) * dilation_rate, stride)
+            in_length = padding + in_length + padding
+
+        if filter_size_int == 1:
+            valid_part = in_length
+        elif isinstance(in_length, Dim):
+            filter_left_dilated = (filter_size_ - 1) * dilation_rate // 2
+            filter_right_dilated = (filter_size_ - 1) * dilation_rate - filter_left_dilated
+            valid_part = in_length.sub_left(filter_left_dilated).sub_right(filter_right_dilated)
+        else:
+            valid_part = in_length - (filter_size_ - 1) * dilation_rate
+
+        if isinstance(valid_part, Dim):
+            out_length = valid_part.ceildiv_right(stride)
+        else:
+            out_length = _ceildiv(valid_part, stride)
+
     else:
         raise ValueError(f"invalid padding {padding!r} (type {type(padding).__name__})")
 
+    if isinstance(in_length, Dim):
+        assert isinstance(out_length, Dim)
+        if name and out_length != in_length:
+            out_length.name = name
+        if in_length.dyn_size_ext is not None and out_length.dyn_size_ext is None:
+            out_dyn_size_ext = calc_conv_out_length(
+                in_length=in_length.dyn_size_ext,
+                filter_size=filter_size,
+                stride=stride,
+                dilation_rate=dilation_rate,
+                padding=padding,
+            )
+            assert isinstance(out_dyn_size_ext, Tensor)
+            out_length.dyn_size_ext = out_dyn_size_ext
+
+    return out_length
+
+
+def make_transposed_conv_out_spatial_dims(
+    in_spatial_dims: Sequence[Dim],
+    *,
+    filter_size: Union[Sequence[Union[int, Dim]], int, Dim],
+    padding: Union[str, int, Sequence[int]],
+    output_padding: Optional[Union[Sequence[Optional[int]], int]] = None,
+    strides: Union[Sequence[Optional[int]], None, int] = None,
+    dilation_rate: Union[Sequence[int], int] = 1,
+    description_prefix: Optional[str] = None,
+) -> List[Dim]:
+    """create out spatial dims from in spatial dims"""
+    nd = len(in_spatial_dims)
+    if isinstance(filter_size, (int, Dim)):
+        filter_size = [filter_size] * nd
+    filter_size = [d.dimension if isinstance(d, Dim) else d for d in filter_size]
+    assert all(isinstance(s, int) for s in filter_size)
+    if isinstance(strides, int) or strides is None:
+        strides = [strides] * nd
+    if isinstance(dilation_rate, int):
+        dilation_rate = [dilation_rate] * nd
+    if isinstance(padding, (int, str)):
+        padding = [padding] * nd
+    if isinstance(output_padding, int) or output_padding is None:
+        output_padding = [output_padding] * nd
+    assert (
+        nd
+        == len(in_spatial_dims)
+        == len(filter_size)
+        == len(strides)
+        == len(dilation_rate)
+        == len(padding)
+        == len(output_padding)
+    )
+    padding = [p.lower() if isinstance(p, str) else p for p in padding]
+    out_spatial_dims = []
+    for i in range(nd):
+        out_spatial_dims.append(
+            calc_transposed_conv_out_length(
+                in_spatial_dims[i],
+                filter_size=filter_size[i],
+                padding=padding[i],
+                stride=strides[i],
+                dilation_rate=dilation_rate[i],
+                name=f"{description_prefix}:spatial{i}" if description_prefix else None,
+            )
+        )
+    return out_spatial_dims
+
+
+def calc_transposed_conv_out_length(
+    in_length: Union[T, int, Dim, Tensor],
+    *,
+    filter_size: Union[int, Dim],
+    padding: Union[int, str],
+    output_padding: Optional[int] = None,
+    stride: Optional[int] = None,
+    dilation_rate: int = 1,
+    name: Optional[str] = None,
+) -> T:
+    """
+    Determines output length of a transposed convolution given input length.
+
+    Copied from TF/Keras conv_utils.deconv_output_length
+    (https://github.com/tensorflow/tensorflow/blob/5912f51d580551e5cee2cfde4cb882594b4d3e60/tensorflow/python/keras/utils/conv_utils.py#L140),
+    adapted with simplification.
+
+    Also see :func:`calc_conv_out_length`.
+
+    :param in_length:
+    :param filter_size:
+    :param padding: one of `"same"`, `"valid"`, `"full"`.
+    :param output_padding: amount of padding along the output dimension.
+        Can be set to `None` in which case the output length is inferred.
+    :param stride:
+    :param dilation_rate:
+    :param name:
+    :returns: The output length (integer)
+    """
+    assert padding in {"same", "valid", "full"} or isinstance(padding, int)
+
+    if isinstance(filter_size, int):
+        filter_size_int = filter_size
+    elif isinstance(filter_size, Dim):
+        filter_size_int = filter_size.dimension
+    else:
+        filter_size_int = None
+    filter_size_ = filter_size_int if isinstance(filter_size_int, int) else filter_size
+
+    # Get the dilated kernel size
+    if dilation_rate != 1 and filter_size_int != 1:
+        filter_size = filter_size + (filter_size_ - 1) * (dilation_rate - 1)
+
+    if stride is None:
+        assert filter_size_int is not None
+        stride = filter_size_int
+    if stride != 1:
+        in_length = in_length * stride
+
+    # Infer length if output padding is None, else compute the exact length
+    if output_padding is None:
+        if padding == "valid" or padding == 0:
+            if filter_size_int == stride:
+                out_length = in_length
+            elif filter_size_int is not None:
+                out_length = in_length + max(filter_size_int - stride, 0)
+            elif isinstance(filter_size, Tensor):
+                out_length = in_length + rf.relu(filter_size - stride)
+            elif isinstance(filter_size, Dim):
+                out_length = in_length + (filter_size - stride)
+            else:
+                raise ValueError(f"invalid filter_size {filter_size!r} type {type(filter_size)}")
+        elif padding == "full":
+            out_length = in_length - (stride + filter_size_ - 2)
+        elif padding == "same":
+            out_length = in_length
+        else:
+            raise ValueError(f"invalid padding {padding!r}")
+
+    else:  # output_padding
+        if padding == "same":
+            pad = filter_size // 2
+        elif padding == "valid":
+            pad = 0
+        elif padding == "full":
+            pad = filter_size - 1
+        elif isinstance(padding, int):
+            pad = padding
+        else:
+            raise ValueError(f"invalid padding {padding!r}")
+        out_length = in_length + (filter_size - stride - 2 * pad + output_padding)
+
+    if isinstance(in_length, Dim):
+        assert isinstance(out_length, Dim)
+        if name and out_length != in_length:
+            out_length.name = name
+        if in_length.dyn_size_ext is not None and out_length.dyn_size_ext is None:
+            out_dyn_size_ext = calc_transposed_conv_out_length(
+                in_length=in_length.dyn_size_ext,
+                filter_size=filter_size,
+                padding=padding,
+                output_padding=output_padding,
+                stride=stride,
+                dilation_rate=dilation_rate,
+            )
+            assert isinstance(out_dyn_size_ext, Tensor)
+            out_length.dyn_size_ext = out_dyn_size_ext
+
+    return out_length
+
+
+def _ceildiv(a: T, b: Union[T, int, Tensor]) -> T:
+    if isinstance(b, int) and b == 1:
+        return a
+    if isinstance(a, Tensor):
+        return rf.ceil_divide(a, b)
+    return -(-a // b)
+
 
 def _should_use_consistent_same_padding() -> bool:
     """

returnn/frontend/decoder/transformer.py

@@ -49,6 +49,7 @@ class TransformerDecoder(rf.Module):
         layer_opts: Optional[Dict[str, Any]] = None,
         embed_dim: Optional[Dim] = None,
         share_embedding: bool = None,
+        input_embedding: bool = True,
         input_embedding_scale: float = None,
         input_dropout: float = None,
         logits_with_bias: bool = False,
@@ -72,6 +73,7 @@ class TransformerDecoder(rf.Module):
         :param layer_opts: options for the decoder layer
         :param embed_dim: if given, will first have an embedding [vocab,embed] and then a linear [embed,model].
         :param share_embedding:
+        :param input_embedding: whether to use input embedding. If False, you must provide input of dimension model_dim.
         :param input_embedding_scale:
         :param input_dropout:
         :param logits_with_bias:
@@ -103,7 +105,7 @@ class TransformerDecoder(rf.Module):
 
         # We could make this optional or configurable if we ever need to.
         # Or maybe you would just have another separate implementation of this module then...
-        self.input_embedding = rf.Embedding(vocab_dim, embed_dim or model_dim)
+        self.input_embedding = rf.Embedding(vocab_dim, embed_dim or model_dim) if input_embedding else None
 
         self.input_embedding_proj = None
         if embed_dim:
@@ -121,21 +123,31 @@ class TransformerDecoder(rf.Module):
             raise TypeError(f"unexpected pos_enc type {pos_enc!r}")
         self.pos_enc = pos_enc
         if share_embedding is None:
-            if BehaviorVersion.get() < 20:
-                logging.getLogger("returnn.frontend").warning(
-                    "TransformerDecoder share_embedding default is False"
-                    f" with your behavior version {BehaviorVersion.get()}."
-                    " Explicitly set share_embedding or switch to a new behavior version >= 20."
-                )
-            share_embedding = True if BehaviorVersion.get() >= 20 else False
+            if embed_dim and embed_dim != model_dim:
+                share_embedding = False
+            elif input_embedding:
+                if BehaviorVersion.get() < 20:
+                    logging.getLogger("returnn.frontend").warning(
+                        "TransformerDecoder share_embedding default is False"
+                        f" with your behavior version {BehaviorVersion.get()}."
+                        " Explicitly set share_embedding or switch to a new behavior version >= 20."
+                    )
+                share_embedding = True if BehaviorVersion.get() >= 20 else False
+            else:  # not input_embedding
+                share_embedding = False
         if input_embedding_scale is None:
-            if BehaviorVersion.get() < 20:
-                logging.getLogger("returnn.frontend").warning(
-                    "TransformerDecoder input_embedding_scale default is suboptimal"
-                    f" with your behavior version {BehaviorVersion.get()}."
-                    " Explicitly set input_embedding_scale or switch to a new behavior version >= 20."
-                )
-            input_embedding_scale = model_dim.dimension**0.5 if BehaviorVersion.get() >= 20 else 1.0
+            if input_embedding:
+                if BehaviorVersion.get() < 20:
+                    logging.getLogger("returnn.frontend").warning(
+                        "TransformerDecoder input_embedding_scale default is suboptimal"
+                        f" with your behavior version {BehaviorVersion.get()}."
+                        " Explicitly set input_embedding_scale or switch to a new behavior version >= 20."
+                    )
+                input_embedding_scale = model_dim.dimension**0.5 if BehaviorVersion.get() >= 20 else 1.0
+            elif pos_enc:
+                input_embedding_scale = model_dim.dimension**0.5
+            else:
+                input_embedding_scale = 1.0
         self.input_embedding_scale = input_embedding_scale
         if input_dropout is None:
             if dropout > 0 and BehaviorVersion.get() < 20:
@@ -179,7 +191,9 @@ class TransformerDecoder(rf.Module):
         self.logits = rf.Linear(model_dim, vocab_dim, with_bias=logits_with_bias)
 
         if share_embedding:
-            assert not embed_dim and not logits_with_bias, "not supported together with share_embedding"
+            assert input_embedding, "input_embedding=True required for share_embedding"
+            assert not embed_dim or embed_dim == model_dim, f"{embed_dim=} not supported with share_embedding"
+            assert not logits_with_bias, "logits_with_bias=True expected with share_embedding"
             self.logits.weight = self.input_embedding.weight
 
     def default_initial_state(self, *, batch_dims: Sequence[Dim]) -> rf.State:
@@ -219,7 +233,12 @@ class TransformerDecoder(rf.Module):
         """
         new_state = rf.State()
 
-        decoded = self.input_embedding(source) * self.input_embedding_scale
+        if self.input_embedding is not None:
+            decoded = self.input_embedding(source)
+        else:
+            decoded = source
+        if self.input_embedding_scale != 1:
+            decoded = decoded * self.input_embedding_scale
         if self.pos_enc is not None:
             decoded = decoded + self.pos_enc(spatial_dim=spatial_dim, offset=state.pos)
         decoded = rf.dropout(decoded, self.input_dropout)

returnn/frontend/encoder/conformer.py

@@ -273,6 +273,7 @@ class ConformerEncoderLayer(rf.Module):
         x_mhsa = self.self_att(x_mhsa_ln, axis=spatial_dim)
         x_mhsa = rf.dropout(x_mhsa, self.dropout, axis=self.dropout_broadcast and self.out_dim)
         x_mhsa_out = x_mhsa + x_ffn1_out
+        del x_mhsa
 
         # Conv
         x_conv_ln = self.conv_layer_norm(x_mhsa_out)

returnn/frontend/encoder/transformer.py

@@ -79,6 +79,8 @@ class TransformerEncoder(rf.Module):
         self.model_dim = model_dim
         self.embed_dim = embed_dim
 
+        self.out_dim = self.model_dim  # alias. consistency, compatibility
+
         if input_embedding is None or isinstance(input_embedding, rf.Module):
             pass
         elif isinstance(input_embedding, type):

returnn/frontend/loss.py

@@ -3,11 +3,12 @@ Loss functions
 """
 
 from __future__ import annotations
+from typing import Optional, Tuple
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 
 
-__all__ = ["cross_entropy", "ctc_loss", "edit_distance"]
+__all__ = ["cross_entropy", "ctc_loss", "ctc_greedy_decode", "edit_distance"]
 
 
 def cross_entropy(
@@ -93,6 +94,44 @@ def ctc_loss(
     )
 
 
+def ctc_greedy_decode(
+    logits: Tensor,
+    *,
+    in_spatial_dim: Dim,
+    blank_index: int,
+    out_spatial_dim: Optional[Dim] = None,
+    target_dim: Optional[Dim] = None,
+    wb_target_dim: Optional[Dim] = None,
+) -> Tuple[Tensor, Dim]:
+    """
+    Greedy CTC decode.
+
+    :return: (labels, out_spatial_dim)
+    """
+    if wb_target_dim is None:
+        assert logits.feature_dim
+        wb_target_dim = logits.feature_dim
+
+    labels = rf.reduce_argmax(logits, axis=wb_target_dim)
+    labels = rf.cast(labels, "int32")
+
+    labels_shifted = rf.shift_right(labels, axis=in_spatial_dim, pad_value=blank_index)
+    mask_repeat = labels != labels_shifted
+    labels, out_spatial_dim = rf.masked_select(
+        labels,
+        mask=(labels != blank_index) & mask_repeat,
+        dims=[in_spatial_dim],
+        out_dim=out_spatial_dim,
+    )
+
+    if target_dim:
+        # Set correct sparse_dim. Only currently implemented if blank comes after.
+        assert target_dim.dimension == blank_index
+        labels.sparse_dim = target_dim
+
+    return labels, out_spatial_dim
+
+
 def edit_distance(a: Tensor, a_spatial_dim: Dim, b: Tensor, b_spatial_dim: Dim, *, dtype: str = "int32") -> Tensor:
     """
     :param a: [B,Ta]

returnn/frontend/module.py

@@ -274,10 +274,17 @@ class Functional(Module):
     (This is often not necessary, but sometimes useful.)
     """
 
-    def __init__(self, func):
+    def __init__(self, func, *, attribs: Optional[Dict[str, Any]] = None):
+        """
+        :param func: callable. you might want to use functools.partial if you want to fix some arguments.
+        :param attribs: optional dict of attributes to set on this module. e.g. ``out_dim``.
+        """
         super().__init__()
         assert callable(func)
         self.func = func
+        if attribs:
+            for k, v in attribs.items():
+                setattr(self, k, v)
 
     def __repr__(self):
         return f"{self.__class__.__name__}({self.func.__qualname__})"

returnn/frontend/nested.py

@@ -275,6 +275,8 @@ def _masked_select(
             return s
         assert s in dim_map
         return dim_map[s]
+    if s is None:
+        return None
     raise TypeError(f"_masked_select: unexpected type ({type(s)})")
 
 
@@ -346,6 +348,7 @@ def _masked_scatter_merge_dims(
     merged_dim_map: Dict[Dim, Dim],
 ) -> T:
     if isinstance(s, Dim):
+        assert isinstance(backup, Dim)
         # This is slightly more complex than in the _masked_select case:
         # We need to merge the s and backup depending on the mask.
         if s in reverse_dim_map:
@@ -353,7 +356,10 @@ def _masked_scatter_merge_dims(
         if s == backup:
             return s
         if s in merged_dim_map:
+            # If this assert fails, see e.g. https://github.com/rwth-i6/returnn/pull/1759 for an example.
+            assert backup in merged_dim_map, f"nested masked_scatter: mismatch of s {s} vs backup {backup}"
             return merged_dim_map[s]
+        assert backup not in merged_dim_map, f"nested masked_scatter: mismatch of s {s} vs backup {backup}"
         # Note: s/backup might even be static dims.
         new_size = _masked_scatter(
             s.get_size_tensor(),
@@ -416,6 +422,9 @@ def _masked_scatter(
         if s in merged_dim_map:
             return merged_dim_map[s]
         return s
+    if s is None:
+        assert backup is None
+        return None
     raise TypeError(f"_masked_scatter: unexpected type ({type(s)})")