Diff Coverage

        Returns:
            Tuple of (adjusted_sparse_mode, adjusted_pre_tokens, adjusted_next_tokens,
                      adjusted_actual_seq_qlen, adjusted_actual_seq_kvlen).
        """
        adjusted_sparse_mode = context.sparse_mode
        adjusted_pre_tokens = context.pre_tokens
        adjusted_next_tokens = context.next_tokens
        adjusted_actual_seq_qlen = context.actual_seq_qlen
        adjusted_actual_seq_kvlen = context.actual_seq_kvlen
        seq_split_num = context.seq_split_num

        if seq_split_num > 1 or context.lb_split_id is not None:
            dynamic_info = self._get_dynamic_shape_info(query, key, input_layout)
            is_dynamic = dynamic_info.get('is_dynamic', False)

            if context.lb_split_id is not None:
                if context.lb_split_num is None:
                    raise ValueError("lb_split_num must not be None when lb_split_id is set")
                split_id = context.lb_split_id
                seq_split_num = context.lb_split_num
            else:
                split_id = self._get_split_id(query_layout, input_layout)
            seq_dim_idx = self._get_seq_dim_idx(self._layout_dims.get(input_layout, {}))

        off_value_layout = off_value.layout if hasattr(off_value, '_layout') else None

        local_args = (indices_local, num_classes, on_value_local, off_value_local, axis)
        cache_values = [indices.layout, on_value_layout, off_value_layout, num_classes, axis]
        return local_args, {}, cache_values

    # pylint: disable=W0237
    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
        """