def scaled_dot_product_attention(context, node):
    """
    Input shapes/types:
    - query : (target_seq, d) or (B, target_seq, d) or (B, h, target_seq, d) or (B,.., target_seq, d)
    - key : (source_seq, d) or (B, source_seq, d) or (B, h, source_seq, d) or (B,.., source_seq, d)
    - value: (source_seq, d_v) or (B, source_seq, d_v) or (B, h, source_seq, d_v) or (B,.., source_seq, d_v)
    - attn_mask : (target_seq, source_seq) or (B, target_seq, source_seq) or (B, h, target_seq, source_seq) or
                  (B, ..., target_seq, source_seq)
    - is_causal : bool
    - scale : optional float

    Output shape: (target_seq, d_v) or (B,...,target_seq, d_v)

    output = softmax(scale*Q*K^transpose + mask) * V

    Currently, Core ML does not support dropout, so it has to be either None or 0

    See details at:
    https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
    """

    def _get_batch_dims(x: Var) -> _utils.SYMBOLIC_SHAPE_TYPE:
        return list(x.shape)[:-2]

    def _broadcast_tensor_to_same_batch_dims(x: Var, batch_dims_var: Var) -> Var:
        x_shape = mb.shape(x=x)
        data_dims_var = mb.slice_by_index(x=x_shape, begin=[x.rank - 2], end=[x.rank])
        broadcast_shape = mb.concat(values=[batch_dims_var, data_dims_var], axis=0)
        broadcast_name = x.name + "_broadcast_to_same_batch_dims"
        return _utils.pymil_broadcast_to(x, broadcast_shape, broadcast_name)

    def _maybe_broadcast_qkv_batch_dims(q: Var, k: Var, v: Var) -> Tuple[Var]:
        q_batch = _get_batch_dims(q)
        k_batch = _get_batch_dims(k)
        v_batch = _get_batch_dims(v)
        symbolic_batch_dims = _utils.pymil_broadcast_shapes([q_batch, k_batch, v_batch])

        q_needs_broadcast = _utils.does_tile_necessary(q_batch, symbolic_batch_dims)
        k_needs_broadcast = _utils.does_tile_necessary(k_batch, symbolic_batch_dims)
        v_needs_broadcast = _utils.does_tile_necessary(v_batch, symbolic_batch_dims)
        if q_needs_broadcast or k_needs_broadcast or v_needs_broadcast:
            batch_dims = _utils.maybe_replace_symbols_with_source_tensor_shape_variables(
                symbolic_batch_dims, [q, k, v]
            )
            batch_dims_var = mb.concat(values=batch_dims, axis=0)
            if q_needs_broadcast:
                q = _broadcast_tensor_to_same_batch_dims(q, batch_dims_var)
            if k_needs_broadcast:
                k = _broadcast_tensor_to_same_batch_dims(k, batch_dims_var)
            if v_needs_broadcast:
                v = _broadcast_tensor_to_same_batch_dims(v, batch_dims_var)
        return q, k, v

    def _parse_positional_args(context, node) -> Tuple[Var]:
        inputs = _get_inputs(context, node, min_expected=3)
        nargs = len(inputs)

        q, k, v = inputs[:3]

        if node.kind == "scaled_dot_product_attention":
            attn_mask = inputs[3] if nargs > 3 else None
            dropout = inputs[4] if nargs > 4 else 0.0
            is_causal = inputs[5].val if nargs > 5 else False
            scale = inputs[6] if nargs > 6 else None
        elif node.kind == "_scaled_dot_product_flash_attention_for_cpu":
            dropout = inputs[3] if nargs > 3 else 0.0
            is_causal = inputs[4].val if nargs > 4 else False
            attn_mask = inputs[5] if nargs > 5 else None
            scale = inputs[6] if nargs > 6 else None
        else:
            assert node.kind in ("coreml.sdpa", "coreml::sdpa")
            attn_mask = inputs[3] if nargs > 3 else None
            dropout = 0.0
            is_causal = False
            scale = None

        return q, k, v, attn_mask, dropout, is_causal, scale

    def _parse_keyword_args(context, node, attn_mask, scale) -> Tuple[Var]:
        attn_mask = _get_kwinputs(context, node, "attn_mask", default=[attn_mask])[0]
        scale = _get_kwinputs(context, node, "scale", default=[scale])[0]
        return attn_mask, scale

    def _check_args(q, k, v, attn_mask, dropout, is_causal) -> None:
        if attn_mask is not None and is_causal:
            raise ValueError(
                "scaled_dot_product_attention op: "
                "attn_mask cannot be provided when is_causal is set to True."
            )

        if dropout is not None:
            if isinstance(dropout, Var):
                if dropout.val is None:
                    raise NotImplementedError(
                        "A variable dropout probability is specified. Since Core ML "
                        "does not support dropout yet, we cowardly refuse to convert it"
                    )
                else:
                    dropout = dropout.val
            if dropout != 0.0:
                raise ValueError(
                    "A non-zero dropout probability is specified. Since Core ML "
                    "does not support dropout yet, we cannot convert it"
                )

        # check that ranks of q, k, v and attn_mask match
        if k.rank != q.rank:
            raise ValueError(
                "Rank of query and key do not match in scaled_dot_product_attention torch op"
            )
        if v.rank != q.rank:
            raise ValueError(
                "Rank of query and value do not match in scaled_dot_product_attention torch op"
            )

    def _construct_causal_attn_mask(query_var: Var, key_var: Var) -> Var:
        # create mask of shape (target_seq, source_seq)
        # s.t the lower triangular of the matrix are all Trues
        # and the strict upper triangular are all Falses
        target_seq = query_var.shape[-2]
        source_seq = key_var.shape[-2]
        if is_symbolic(target_seq) or is_symbolic(source_seq):
            target_seq = _utils.pymil_value_at(mb.shape(x=query_var), -2)
            source_seq = _utils.pymil_value_at(mb.shape(x=key_var), -2)
        mask_shape = mb.concat(values=(target_seq, source_seq), axis=0)
        all_trues = mb.fill(shape=mask_shape, value=True)
        mask = mb.band_part(x=all_trues, lower=-1, upper=0)
        return mask

    def _cast_bool_attn_mask(attn_mask: Var, query_var: Var) -> Var:
        """
        compute float mask as (1 - cast(bool_mask)) * -30k
        """
        assert is_bool(attn_mask.dtype)

        mask = mb.cast(x=attn_mask, dtype=types.builtin_to_string(query_var.dtype))
        compliment_of_mask = mb.sub(x=1.0, y=mask)
        # Use a big enough but not easily fp16-overflow number (e.g. -3e4) as -inf
        float_mask = mb.mul(x=-3e4, y=compliment_of_mask)
        return float_mask

    def _translate_torch_args(q, k, attn_mask, is_causal, can_use_fused_sdpa) -> Var:
        mask = attn_mask
        if is_causal:
            mask = _construct_causal_attn_mask(q, k)
        if mask is not None and is_bool(mask.dtype) and not can_use_fused_sdpa:
            # In Core ML, only fused sdpa can use bool attention mask
            # i.e. the decomposition has to use floating point mask
            mask = _cast_bool_attn_mask(mask, q)
        return mask

    q, k, v, attn_mask, dropout, is_causal, scale = _parse_positional_args(context, node)
    attn_mask, scale = _parse_keyword_args(context, node, attn_mask, scale)
    _check_args(q, k, v, attn_mask, dropout, is_causal)

    # Since ios18, Core ML supports scaled_dot_product_attention op
    # It does not have scale, though
    can_use_fused_sdpa = is_current_opset_version_compatible_with(target.iOS18) and scale is None
    mask = _translate_torch_args(q, k, attn_mask, is_causal, can_use_fused_sdpa)

    if can_use_fused_sdpa:
        # ios18 scaled_dot_product_attention only supports rank >= 3
        is_rank_2 = q.rank == 2

        if is_rank_2:
            q = mb.expand_dims(x=q, axes=[0])
            k = mb.expand_dims(x=k, axes=[0])
            v = mb.expand_dims(x=v, axes=[0])

        q, k, v = _maybe_broadcast_qkv_batch_dims(q, k, v)
        res = mb.scaled_dot_product_attention(
            query=q, key=k, value=v, attn_mask=mask, name=node.name
        )

        if is_rank_2:
            res = mb.squeeze(x=res, axes=[0], name=node.name)

    else:
        res = _utils._decompose_scaled_dot_product_attention(q, k, v, mask, node.name, scale=scale)

    context.add(res)

    Input shapes/types:
    - query : (target_seq, d) or (B, target_seq, d) or (B, h, target_seq, d) or (B,.., target_seq, d)
    - key : (source_seq, d) or (B, source_seq, d) or (B, h, source_seq, d) or (B,.., source_seq, d)
    - value: (source_seq, d_v) or (B, source_seq, d_v) or (B, h, source_seq, d_v) or (B,.., source_seq, d_v)
    - attn_mask : (target_seq, source_seq) or (B, target_seq, source_seq) or (B, h, target_seq, source_seq) or
                  (B, ..., target_seq, source_seq)
    - is_causal : bool
    - scale : optional float

    Output shape: (target_seq, d_v) or (B,...,target_seq, d_v)

    output = softmax(scale*Q*K^transpose + mask) * V

    Currently, Core ML does not support dropout, so it has to be either None or 0

    See details at:
    https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
    
coremltools/converters/mil/frontend/torch/ops.py