[Multimodal] Support n>1 outputs for GLM-Image generation (#31027)

Co-authored-by: AuFlow <AuFlow@users.noreply.github.com>

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -1183,7 +1183,9 @@ class ImagePipelineConfig(PipelineConfig):
 
         latents = maybe_unpad_latents(latents, batch)
 
-        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.reshape(
+            batch_size, height // 2, width // 2, channels // 4, 2, 2
+        )
         latents = latents.permute(0, 3, 1, 4, 2, 5)
         return latents, batch_size, channels, height, width
 

python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py

@@ -86,12 +86,12 @@ class GlmImagePipelineConfig(SpatialImagePipelineConfig):
     def get_decode_scale_and_shift(self, device, dtype, vae):
         latents_mean = (
             torch.tensor(self.vae_config.latents_mean)
-            .view(1, self.vae_config.latent_channels, 1, 1)
+            .reshape(1, self.vae_config.latent_channels, 1, 1)
             .to(device, dtype)
         )
         latents_std = (
             torch.tensor(self.vae_config.latents_std)
-            .view(1, self.vae_config.latent_channels, 1, 1)
+            .reshape(1, self.vae_config.latent_channels, 1, 1)
             .to(device, dtype)
         )
         return 1.0 / latents_std, latents_mean

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -452,6 +452,35 @@ def _ensure_contiguous(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]
     return tensor.contiguous() if tensor is not None else None
 
 
+def _is_scalar_or_hidden_modulation(tensor: torch.Tensor, hidden_size: int) -> bool:
+    return tensor.numel() in (1, hidden_size)
+
+
+def _can_use_npu_fused_scale_shift(
+    x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
+) -> bool:
+    hidden_size = x.shape[-1]
+    return _is_scalar_or_hidden_modulation(scale, hidden_size) and (
+        _is_scalar_or_hidden_modulation(shift, hidden_size)
+        or tuple(shift.shape) == tuple(x.shape)
+    )
+
+
+def _try_npu_fused_scale_shift(
+    x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
+) -> torch.Tensor | None:
+    if not _can_use_npu_fused_scale_shift(x, shift, scale):
+        return None
+
+    from sgl_kernel_npu.norm.scale_shift import fused_scale_shift
+
+    scale = scale.reshape(-1)
+    if tuple(shift.shape) != tuple(x.shape):
+        shift = shift.reshape(-1)
+
+    return fused_scale_shift(x, scale.contiguous(), shift.contiguous())
+
+
 class _ScaleResidualNormScaleShift(CustomOp):
     """
     Fused kernel that combines:
@@ -607,8 +636,6 @@ class _ScaleResidualNormScaleShift(CustomOp):
         shift: torch.Tensor,
         scale: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        from sgl_kernel_npu.norm.scale_shift import fused_scale_shift
-
         # x.shape: [batch_size, seq_len, inner_dim]
         if isinstance(gate, int):
             # used by cross-attention, should be 1
@@ -628,7 +655,9 @@ class _ScaleResidualNormScaleShift(CustomOp):
         else:
             raise ValueError(f"Gate type {type(gate)} not supported")
         normalized = self.norm(residual_output)
-        modulated = fused_scale_shift(normalized, scale, shift)
+        modulated = _try_npu_fused_scale_shift(normalized, shift, scale)
+        if modulated is None:
+            modulated = normalized * (1 + scale) + shift
         return modulated, residual_output
 
 
@@ -747,23 +776,12 @@ class _NormScaleShift(CustomOp):
     def forward_npu(
         self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
     ) -> torch.Tensor:
-        hidden_size = x.shape[-1]
-        x_numel = x.numel()
-
-        if scale.numel() in (1, hidden_size) and shift.numel() in (
-            1,
-            hidden_size,
-            x_numel,
-        ):
-            from sgl_kernel_npu.norm.scale_shift import fused_scale_shift
-
-            normalized = self.norm(x)
-            modulated = fused_scale_shift(
-                normalized, scale.contiguous(), shift.contiguous()
-            )
+        normalized = self.norm(x)
+        modulated = _try_npu_fused_scale_shift(normalized, shift, scale)
+        if modulated is not None:
             return modulated.to(x.dtype)
 
-        return self.forward_native(x, shift, scale)
+        return (normalized * (1 + scale) + shift).to(x.dtype)
 
 
 class LayerNormScaleShift(_NormScaleShift):

python/sglang/multimodal_gen/runtime/models/dits/glm_image.py

@@ -201,10 +201,10 @@ class GlmImageCombinedTimestepSizeEmbeddings(nn.Module):
     ) -> torch.Tensor:
         timesteps_proj = self.time_proj(timestep)
 
-        crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(
+        crop_coords_proj = self.condition_proj(crop_coords.flatten()).reshape(
             crop_coords.size(0), -1
         )
-        target_size_proj = self.condition_proj(target_size.flatten()).view(
+        target_size_proj = self.condition_proj(target_size.flatten()).reshape(
             target_size.size(0), -1
         )
 
@@ -908,11 +908,50 @@ class GlmImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
 
         batch_size, num_channels, height, width = hidden_states.shape
 
-        timestep = timestep - 1.0
-
         if isinstance(encoder_hidden_states, list):
             encoder_hidden_states = encoder_hidden_states[0]
 
+        if current_platform.is_npu() and batch_size > 1 and kv_caches is None:
+
+            def slice_batch(value, index):
+                if (
+                    isinstance(value, torch.Tensor)
+                    and value.dim() > 0
+                    and value.shape[0] == batch_size
+                ):
+                    return value[index : index + 1]
+                return value
+
+            outputs = []
+            for index in range(batch_size):
+                sample_encoder_hidden_states = encoder_hidden_states[index : index + 1]
+                sample_attention_mask = slice_batch(attention_mask, index)
+                if isinstance(sample_attention_mask, torch.Tensor):
+                    valid_text_length = int(sample_attention_mask.sum().item())
+                    sample_encoder_hidden_states = sample_encoder_hidden_states[
+                        :, :valid_text_length
+                    ]
+                    sample_attention_mask = None
+
+                outputs.append(
+                    self.forward(
+                        hidden_states=hidden_states[index : index + 1],
+                        encoder_hidden_states=sample_encoder_hidden_states,
+                        prior_token_id=prior_token_id[index : index + 1],
+                        prior_token_drop=prior_token_drop[index : index + 1],
+                        timestep=slice_batch(timestep, index),
+                        target_size=slice_batch(target_size, index),
+                        crop_coords=slice_batch(crop_coords, index),
+                        attention_kwargs=attention_kwargs,
+                        attention_mask=sample_attention_mask,
+                        freqs_cis=freqs_cis,
+                        guidance=slice_batch(guidance, index),
+                    )
+                )
+            return torch.cat(outputs, dim=0)
+
+        timestep = timestep - 1.0
+
         # 1. RoPE
         image_rotary_emb = freqs_cis
         if image_rotary_emb is None:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py

@@ -126,6 +126,33 @@ def pooled_image_features_to_tensor(image_features) -> torch.Tensor:
     return torch.cat(tuple(image_features), dim=0)
 
 
+def _num_outputs_per_prompt(batch: Req) -> int:
+    return max(1, int(getattr(batch, "num_outputs_per_prompt", 1) or 1))
+
+
+def _seed_for_output(seed: Optional[Union[int, List[int]]], output_idx: int):
+    if seed is None:
+        return None
+    if isinstance(seed, list):
+        if not seed:
+            return None
+        if output_idx < len(seed):
+            return int(seed[output_idx])
+        return int(seed[0]) + output_idx
+    return int(seed) + output_idx
+
+
+def _repeat_to_batch(tensor: Optional[torch.Tensor], batch_size: int):
+    if tensor is None or tensor.shape[0] == batch_size:
+        return tensor
+    if tensor.shape[0] != 1:
+        raise ValueError(
+            f"Cannot expand tensor with batch size {tensor.shape[0]} to {batch_size}"
+        )
+    repeat_shape = [batch_size] + [1] * (tensor.dim() - 1)
+    return tensor.repeat(*repeat_shape)
+
+
 class GlmImageAR(PipelineStage):
     r"""
     Pipeline for text-to-image generation using GLM-Image.
@@ -180,11 +207,11 @@ class GlmImageAR(PipelineStage):
     def _upsample_token_ids(
         token_ids: torch.Tensor, token_h: int, token_w: int
     ) -> torch.Tensor:
-        token_ids = token_ids.view(1, 1, token_h, token_w)
+        token_ids = token_ids.reshape(1, 1, token_h, token_w)
         token_ids = torch.nn.functional.interpolate(
             token_ids.float(), scale_factor=2, mode="nearest"
         ).to(dtype=torch.long)
-        token_ids = token_ids.view(1, -1)
+        token_ids = token_ids.reshape(1, -1)
         return token_ids
 
     def generate_prior_tokens(
@@ -386,28 +413,51 @@ class GlmImageAR(PipelineStage):
             width = width or ar_condition_images[0].width
 
         time_start = time.time()
+        num_outputs = _num_outputs_per_prompt(batch)
         seed = getattr(batch, "seed", None)
-        if seed is None:
-            prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
-                prompt=prompt,
-                image=ar_condition_images,
-                height=height,
-                width=width,
-                server_args=server_args,
-            )
-        else:
-            rng_devices = []
-            if device.type == "cuda":
-                rng_devices.append(torch.cuda.current_device())
-            with torch.random.fork_rng(devices=rng_devices, enabled=True):
-                torch.manual_seed(int(seed))
-                prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
-                    prompt=prompt,
-                    image=ar_condition_images,
-                    height=height,
-                    width=width,
-                    server_args=server_args,
+        rng_devices = []
+        rng_device_type = "cuda"
+        if device.type == "cuda":
+            rng_devices.append(torch.cuda.current_device())
+        elif device.type == "npu":
+            rng_devices.append(torch.npu.current_device())
+            rng_device_type = "npu"
+
+        prior_token_ids = []
+        prior_token_image_ids = None
+        for output_idx in range(num_outputs):
+            output_seed = _seed_for_output(seed, output_idx)
+            if output_seed is None:
+                prior_token_id, output_prior_token_image_ids = (
+                    self.generate_prior_tokens(
+                        prompt=prompt,
+                        image=ar_condition_images,
+                        height=height,
+                        width=width,
+                        server_args=server_args,
+                    )
                 )
+            else:
+                with torch.random.fork_rng(
+                    devices=rng_devices,
+                    enabled=True,
+                    device_type=rng_device_type,
+                ):
+                    torch.manual_seed(output_seed)
+                    prior_token_id, output_prior_token_image_ids = (
+                        self.generate_prior_tokens(
+                            prompt=prompt,
+                            image=ar_condition_images,
+                            height=height,
+                            width=width,
+                            server_args=server_args,
+                        )
+                    )
+            prior_token_ids.append(prior_token_id)
+            if prior_token_image_ids is None:
+                prior_token_image_ids = output_prior_token_image_ids
+
+        prior_token_id = torch.cat(prior_token_ids, dim=0)
         prior_token_id = prior_token_id.to(device=device)
         time_end = time.time()
         logger.info(f"generate_prior_tokens time: {time_end - time_start}")
@@ -579,7 +629,7 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
 
         seq_len = prompt_embeds.size(1)
         prompt_embeds = prompt_embeds.repeat(1, 1, 1)
-        prompt_embeds = prompt_embeds.view(1, seq_len, -1)
+        prompt_embeds = prompt_embeds.reshape(1, seq_len, -1)
 
         negative_prompt_embeds = None
         if do_classifier_free_guidance:
@@ -608,7 +658,7 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
 
             seq_len = negative_prompt_embeds.size(1)
             negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(1, seq_len, -1)
+            negative_prompt_embeds = negative_prompt_embeds.reshape(1, seq_len, -1)
 
         return prompt_embeds, negative_prompt_embeds
 
@@ -726,8 +776,28 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         width = batch.width
 
         device = get_local_torch_device()
+        batch_size = _num_outputs_per_prompt(batch)
         max_sequence_length = 1024
-        generator = torch.Generator(device=device).manual_seed(batch.seed)
+        seed = getattr(batch, "seed", None)
+        if batch_size == 1:
+            output_seed = _seed_for_output(seed, 0)
+            generator = (
+                None
+                if output_seed is None
+                else torch.Generator(device=device).manual_seed(int(output_seed))
+            )
+        elif seed is None:
+            generator = None
+        else:
+            output_seeds = [_seed_for_output(seed, i) for i in range(batch_size)]
+            generator = (
+                None
+                if any(output_seed is None for output_seed in output_seeds)
+                else [
+                    torch.Generator(device=device).manual_seed(int(output_seed))
+                    for output_seed in output_seeds
+                ]
+            )
         attention_kwargs = {}
         prompt_embeds = None
         do_classifier_free_guidance = True
@@ -744,6 +814,7 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         prior_token_id = batch.prior_token_id
         prior_token_image_ids = batch.prior_token_image_ids
         prior_token_id = prior_token_id.to(device)
+        prior_token_id = _repeat_to_batch(prior_token_id, batch_size)
 
         # 3. Encode input prompt
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
@@ -754,6 +825,8 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
             device=device,
             dtype=dtype,
         )
+        prompt_embeds = _repeat_to_batch(prompt_embeds, batch_size)
+        negative_prompt_embeds = _repeat_to_batch(negative_prompt_embeds, batch_size)
 
         # 4. process images
         if ar_condition_images is not None:
@@ -776,7 +849,7 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         # 5. Prepare latents and (optional) condition_images kv cache
         latent_channels = self.transformer.config.in_channels
         latents = self.prepare_latents(
-            batch_size=1,
+            batch_size=batch_size,
             num_channels_latents=latent_channels,
             height=height,
             width=width,
@@ -788,10 +861,10 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         kv_caches = GlmImageKVCache(num_layers=self.transformer.config.num_layers)
 
         if ar_condition_images is not None:
-            latents_mean = torch.tensor(self.vae.config.latents_mean).view(
+            latents_mean = torch.tensor(self.vae.config.latents_mean).reshape(
                 1, self.vae.config.latent_channels, 1, 1
             )
-            latents_std = torch.tensor(self.vae.config.latents_std).view(
+            latents_std = torch.tensor(self.vae.config.latents_std).reshape(
                 1, self.vae.config.latent_channels, 1, 1
             )
 
@@ -805,6 +878,7 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
                 condition_image = align_tensor_to_module_dtype(
                     condition_image, self.vae, device=device
                 )
+                condition_image = _repeat_to_batch(condition_image, batch_size)
 
                 condition_latent = retrieve_latents(
                     self.vae.encode(condition_image),
@@ -812,6 +886,13 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
                     sample_mode="argmax",
                 )
                 condition_latent = (condition_latent - latents_mean) / latents_std
+                if condition_image_prior_token_id.dim() == 1:
+                    condition_image_prior_token_id = (
+                        condition_image_prior_token_id.unsqueeze(0)
+                    )
+                condition_image_prior_token_id = _repeat_to_batch(
+                    condition_image_prior_token_id.to(device=device), batch_size
+                )
 
                 # Do not remove.
                 # It would be use to run the reference image through a
@@ -827,17 +908,23 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
                         _ = transformer(
                             hidden_states=condition_latent,
                             encoder_hidden_states=torch.zeros_like(prompt_embeds)[
-                                :1, :0, ...
+                                :, :0, ...
                             ],
                             prior_token_id=condition_image_prior_token_id,
                             prior_token_drop=torch.full_like(
                                 condition_image_prior_token_id, False, dtype=torch.bool
                             ),
-                            timestep=torch.zeros((1,), device=device),
+                            timestep=torch.zeros((batch_size,), device=device),
                             target_size=torch.tensor(
-                                [condition_image.shape[-2:]], device=device
+                                [condition_image.shape[-2:]],
+                                dtype=prompt_embeds.dtype,
+                                device=device,
+                            ).repeat(batch_size, 1),
+                            crop_coords=torch.zeros(
+                                (batch_size, 2),
+                                dtype=prompt_embeds.dtype,
+                                device=device,
                             ),
-                            crop_coords=torch.zeros((1, 2), device=device),
                             attention_kwargs=attention_kwargs,
                             kv_caches=kv_caches,
                             kv_caches_mode="write",
@@ -847,10 +934,10 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         target_size = (height, width)
         target_size = torch.tensor(
             [target_size], dtype=prompt_embeds.dtype, device=device
-        )
+        ).repeat(batch_size, 1)
         crops_coords_top_left = torch.tensor(
             [(0, 0)], dtype=prompt_embeds.dtype, device=device
-        )
+        ).repeat(batch_size, 1)
 
         # Prepare timesteps
         scheduler = self.scheduler

python/sglang/multimodal_gen/test/unit/test_glm_image_multi_output.py

@@ -0,0 +1,170 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    _can_use_npu_fused_scale_shift,
+)
+from sglang.multimodal_gen.runtime.models.dits import glm_image as glm_model
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages import (
+    glm_image as glm_stage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import (
+    GlmImageAR,
+    GlmImageBeforeDenoisingStage,
+)
+
+
+class _DummyVisionLanguageEncoder:
+    device = torch.device("cpu")
+
+
+class _RecordingGlmImageAR(GlmImageAR):
+    def __init__(self):
+        super().__init__(
+            processor=None, vision_language_encoder=_DummyVisionLanguageEncoder()
+        )
+        self.initial_seeds = []
+
+    def generate_prior_tokens(self, **kwargs):
+        self.initial_seeds.append(torch.initial_seed())
+        output_idx = len(self.initial_seeds)
+        return torch.full((1, 4), output_idx, dtype=torch.long), None
+
+
+class _DummySchedulerConfig(dict):
+    num_train_timesteps = 1000
+
+
+class _DummyScheduler:
+    config = _DummySchedulerConfig(
+        base_image_seq_len=256,
+        base_shift=0.25,
+        max_shift=0.75,
+    )
+
+    def set_timesteps(self, timesteps=None, sigmas=None, device=None, **kwargs):
+        self.timesteps = torch.as_tensor(timesteps, device=device)
+
+
+class _RecordingBeforeDenoisingStage(GlmImageBeforeDenoisingStage):
+    def __init__(self):
+        self.transformer = SimpleNamespace(
+            config=SimpleNamespace(
+                in_channels=4,
+                num_layers=1,
+                patch_size=2,
+            )
+        )
+        self.vae = SimpleNamespace(config=SimpleNamespace(block_out_channels=[1]))
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.scheduler = _DummyScheduler()
+
+    def encode_prompt(self, *args, **kwargs):
+        return torch.ones(1, 3, 5), torch.zeros(1, 3, 5)
+
+    def prepare_latents(
+        self, batch_size, num_channels_latents, height, width, **kwargs
+    ):
+        return torch.zeros(batch_size, num_channels_latents, height // 2, width // 2)
+
+
+class _NpuForwardProbe(glm_model.GlmImageTransformer2DModel):
+    def __init__(self):
+        pass
+
+    def forward(self, **kwargs):
+        batch_size = kwargs["hidden_states"].shape[0]
+        if batch_size > 1:
+            return super().forward(**kwargs)
+        marker = kwargs["prior_token_id"].float().view(1, 1, 1, 1)
+        return marker.expand(1, 1, 2, 2)
+
+
+def test_npu_fused_scale_shift_accepts_broadcast_modulation_shapes():
+    hidden_size = 8
+    x = torch.zeros(1, 4, hidden_size)
+
+    for shape in ((), (1,), (hidden_size,), (1, hidden_size), (1, 1, hidden_size)):
+        scale = torch.zeros(shape)
+        shift = torch.zeros(shape)
+        assert _can_use_npu_fused_scale_shift(x, shift, scale)
+
+    assert not _can_use_npu_fused_scale_shift(
+        x, torch.zeros(1, 2, hidden_size), torch.zeros(1, 2, hidden_size)
+    )
+
+
+def test_ar_stage_generates_one_prior_per_requested_output():
+    stage = _RecordingGlmImageAR()
+    batch = SimpleNamespace(
+        prompt="a cat",
+        height=64,
+        width=64,
+        image_path=None,
+        num_outputs_per_prompt=2,
+        seed=11,
+    )
+
+    with patch.object(
+        glm_stage, "get_local_torch_device", return_value=torch.device("cpu")
+    ):
+        result = stage.forward(batch, SimpleNamespace())
+
+    assert result.prior_token_id.shape == (2, 4)
+    assert result.prior_token_id.tolist() == [[1, 1, 1, 1], [2, 2, 2, 2]]
+    assert stage.initial_seeds == [11, 12]
+
+
+def test_before_denoising_expands_latents_and_conditions_for_requested_outputs():
+    stage = _RecordingBeforeDenoisingStage()
+    batch = SimpleNamespace(
+        prompt="a cat",
+        height=64,
+        width=64,
+        image_path=None,
+        guidance_scale=4.5,
+        num_inference_steps=2,
+        num_outputs_per_prompt=2,
+        seed=7,
+        prior_token_id=torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]),
+        prior_token_image_ids=None,
+    )
+
+    with patch.object(
+        glm_stage, "get_local_torch_device", return_value=torch.device("cpu")
+    ):
+        result = stage.forward(batch, SimpleNamespace())
+
+    assert result.latents.shape[0] == 2
+    assert result.prompt_embeds[0].shape[0] == 2
+    assert result.negative_prompt_embeds[0].shape[0] == 2
+    assert result.prior_token_id.shape == (2, 4)
+    assert result.prior_token_drop_cond.shape == (2, 4)
+    assert result.target_size.shape == (2, 2)
+    assert result.crop_coords.shape == (2, 2)
+
+
+def test_npu_transformer_fallback_runs_each_batch_item_independently():
+    model = _NpuForwardProbe()
+    hidden_states = torch.zeros(2, 1, 2, 2)
+    encoder_hidden_states = torch.zeros(2, 3, 4)
+    prior_token_id = torch.tensor([[3], [9]])
+    attention_mask = torch.tensor([[True, True, False], [True, False, False]])
+
+    with patch.object(glm_model.current_platform, "is_npu", return_value=True):
+        output = glm_model.GlmImageTransformer2DModel.forward(
+            model,
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            prior_token_id=prior_token_id,
+            prior_token_drop=torch.zeros_like(prior_token_id, dtype=torch.bool),
+            timestep=torch.ones(2),
+            target_size=torch.ones(2, 2),
+            crop_coords=torch.zeros(2, 2),
+            attention_mask=attention_mask,
+        )
+
+    assert output.shape == (2, 1, 2, 2)
+    assert output[:, 0, 0, 0].tolist() == [3.0, 9.0]