def test_different_batch_dims(self, compute_unit, backend, frontend, minimum_deployment_target):
        """
        The query/key/value inputs can have different batch_dims.
        """
        q_shape = [1, 2, 10, 3]
        k_shape = [2, 1, 10, 3]
        v_shape = [2, 2, 10, 3]
        input_shape = [
            q_shape,
            k_shape,
            v_shape,
        ]

        model = ModuleWrapper(
            function=nn.functional.scaled_dot_product_attention,
            kwargs={
                "attn_mask": None,
                "dropout_p": 0.0,
                "is_causal": False,
            },
        )

        res = self.run_compare_torch(
            input_shape,
            model,
            frontend=frontend,
            backend=backend,
            compute_unit=compute_unit,
            minimum_deployment_target=minimum_deployment_target,
        )

        # mb.sdpa is introduced in iOS 18, so before iOS 18 we would decompose sdpa
        # torch.sdpa is not a core aten op, so executorch would decompose sdpa
        if (
            backend[0] == "mlprogram"
            and minimum_deployment_target == ct.target.iOS18
            and frontend != TorchFrontend.EXECUTORCH
        ):
            if backend[1] == "fp16":
                expected_ops = [
                    "cast",
                    "tile",
                    "cast",
                    "tile",
                    "cast",
                    "scaled_dot_product_attention",
                ]
            else:
                expected_ops = ["tile", "tile", "scaled_dot_product_attention"]
            assert get_op_types_in_program(res[1]._mil_program) == expected_ops

        The query/key/value inputs can have different batch_dims.
        
coremltools/converters/mil/frontend/torch/test/test_torch_ops.py