code: def check_use_cache_forward_with_attn_mask(
      self, model_class_name, config, inputs_dict
  ):
      max_decoder_length = 20
      model = model_class_name(config)
  
      encoder_outputs = model.encode(inputs_dict["input_ids"])
  
      decoder_input_ids, decoder_attention_mask = (
          inputs_dict["decoder_input_ids"],
          inputs_dict["decoder_attention_mask"],
      )
  
      decoder_attention_mask_cache = jnp.concatenate(
          [
              decoder_attention_mask,
              jnp.zeros(
                  (
                      decoder_attention_mask.shape[0],
                      max_decoder_length
                      - decoder_attention_mask.shape[1],
                  )
              ),
          ],
          axis=-1,
      )
  
      past_key_values = model.init_cache(
          decoder_input_ids.shape[0],
          max_decoder_length,
          encoder_outputs,
      )
      decoder_position_ids = jnp.broadcast_to(
          jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :],
          (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1),
      )
  
      outputs_cache = model.decode(
          decoder_input_ids[:, :-1],
          encoder_outputs,
          decoder_attention_mask=decoder_attention_mask_cache,
          past_key_values=past_key_values,
          decoder_position_ids=decoder_position_ids,
      )
      decoder_position_ids = jnp.array(
          decoder_input_ids.shape[0]
          * [[decoder_input_ids.shape[-1] - 1]],
          dtype="i4",
      )
      outputs_cache_next = model.decode(
          decoder_input_ids[:, -1:],
          encoder_outputs,
          past_key_values=outputs_cache.past_key_values,
          decoder_attention_mask=decoder_attention_mask_cache,
          decoder_position_ids=decoder_position_ids,
      )
  
      outputs = model.decode(
          decoder_input_ids,
          encoder_outputs,
          decoder_attention_mask=decoder_attention_mask,
      )
  
      diff = np.max(
          np.abs(
              (outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])
          )
      )
      self.parent.assertTrue(diff < 1e-3, msg=f"Max diff is {diff}")
  
imports: import unittest
  import numpy
  import timeout_decorator
  from transformers import MBartConfig, is_flax_available
  from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
  from transformers.utils import cached_property
  from ...generation.test_generation_flax_utils import FlaxGenerationTesterMixin
  from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor
label_desc: Use image node_id 5 for calling the FlaxMBartModelTester obj's underlying member method code with example usage: obj.check_use_cache_forward_with_attn_mask(model_class_name, config, inputs_dict) without return types