=================================== FAILURES ===================================
  __________ TestAttentionValueScaleForward.test_scale_applied_to_value __________
  
  self = <tests.unit_tests.models.mimo_v2_flash.test_modeling_mimo_v2_flash.TestAttentionValueScaleForward object at 0x79d70ac1fd40>
  
      def test_scale_applied_to_value(self):
          scale = 0.707
          v = torch.randn(2, 4, 8, 64)
  >       captured, _, _ = _invoke_forward(scale, v)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^
  
  tests/unit_tests/models/mimo_v2_flash/test_modeling_mimo_v2_flash.py:98: 
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  tests/unit_tests/models/mimo_v2_flash/test_modeling_mimo_v2_flash.py:75: in _invoke_forward
      out = instance.forward(
  src/megatron/bridge/models/mimo_v2_flash/modeling_mimo_v2_flash.py:273: in forward
      return super().forward(query, key, value, attention_mask, attn_mask_type, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  <string>:2: in forward
      ???
  /usr/lib/python3.12/unittest/mock.py:192: in checksig
      sig.bind(*args, **kwargs)
  /usr/lib/python3.12/inspect.py:3242: in bind
      return self._bind(args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  
  self = <Signature (self, query: 'Tensor', key: 'Tensor', value: 'Tensor', attention_mask: 'Optional[Tensor]', attn_mask_type:...' = None, packed_seq_params: 'Optional[PackedSeqParams]' = None, num_splits: 'Optional[int]' = None) -> 'torch.Tensor'>
  args = (<[AttributeError("'MiMoV2FlashTEDotProductAttention' object has no attribute '_modules'") raised in repr()] MiMoV2Fla...e-01,  8.8249e-01,  ..., -7.7397e-02,
             -1.3442e+00,  6.1891e-01]]]]), None, <MagicMock id='133991313938896'>)
  kwargs = {'extra_kwarg': 'passthrough'}
  
      def _bind(self, args, kwargs, *, partial=False):
          """Private method. Don't use directly."""
      
          arguments = {}
      
          parameters = iter(self.parameters.values())
          parameters_ex = ()
          arg_vals = iter(args)
      
          while True:
              # Let's iterate through the positional arguments and corresponding
              # parameters
              try:
                  arg_val = next(arg_vals)
              except StopIteration:
                  # No more positional arguments
                  try:
                      param = next(parameters)
                  except StopIteration:
                      # No more parameters. That's it. Just need to check that
                      # we have no `kwargs` after this while loop
                      break
                  else:
                      if param.kind == _VAR_POSITIONAL:
                          # That's OK, just empty *args.  Let's start parsing
                          # kwargs
                          break
                      elif param.name in kwargs:
                          if param.kind == _POSITIONAL_ONLY:
                              msg = '{arg!r} parameter is positional only, ' \
                                    'but was passed as a keyword'
                              msg = msg.format(arg=param.name)
                              raise TypeError(msg) from None
                          parameters_ex = (param,)
                          break
                      elif (param.kind == _VAR_KEYWORD or
                                                  param.default is not _empty):
                          # That's fine too - we have a default value for this
                          # parameter.  So, lets start parsing `kwargs`, starting
                          # with the current parameter
                          parameters_ex = (param,)
                          break
                      else:
                          # No default, not VAR_KEYWORD, not VAR_POSITIONAL,
                          # not in `kwargs`
                          if partial:
                              parameters_ex = (param,)
                              break
                          else:
                              if param.kind == _KEYWORD_ONLY:
                                  argtype = ' keyword-only'
                              else:
                                  argtype = ''
                              msg = 'missing a required{argtype} argument: {arg!r}'
                              msg = msg.format(arg=param.name, argtype=argtype)
                              raise TypeError(msg) from None
              else:
                  # We have a positional argument to process
                  try:
                      param = next(parameters)
                  except StopIteration:
                      raise TypeError('too many positional arguments') from None
                  else:
                      if param.kind in (_VAR_KEYWORD, _KEYWORD_ONLY):
                          # Looks like we have no parameter for this positional
                          # argument
                          raise TypeError(
                              'too many positional arguments') from None
      
                      if param.kind == _VAR_POSITIONAL:
                          # We have an '*args'-like argument, let's fill it with
                          # all positional arguments we have left and move on to
                          # the next phase
                          values = [arg_val]
                          values.extend(arg_vals)
                          arguments[param.name] = tuple(values)
                          break
      
                      if param.name in kwargs and param.kind != _POSITIONAL_ONLY:
                          raise TypeError(
                              'multiple values for argument {arg!r}'.format(
                                  arg=param.name)) from None
      
                      arguments[param.name] = arg_val
      
          # Now, we iterate through the remaining parameters to process
          # keyword arguments
          kwargs_param = None
          for param in itertools.chain(parameters_ex, parameters):
              if param.kind == _VAR_KEYWORD:
                  # Memorize that we have a '**kwargs'-like parameter
                  kwargs_param = param
                  continue
      
              if param.kind == _VAR_POSITIONAL:
                  # Named arguments don't refer to '*args'-like parameters.
                  # We only arrive here if the positional arguments ended
                  # before reaching the last parameter before *args.
                  continue
      
              param_name = param.name
              try:
                  arg_val = kwargs.pop(param_name)
              except KeyError:
                  # We have no value for this parameter.  It's fine though,
                  # if it has a default value, or it is an '*args'-like
                  # parameter, left alone by the processing of positional
                  # arguments.
                  if (not partial and param.kind != _VAR_POSITIONAL and
                                                      param.default is _empty):
                      raise TypeError('missing a required argument: {arg!r}'. \
                                      format(arg=param_name)) from None
      
              else:
                  if param.kind == _POSITIONAL_ONLY:
                      # This should never happen in case of a properly built
                      # Signature object (but let's have this check here
                      # to ensure correct behaviour just in case)
                      raise TypeError('{arg!r} parameter is positional only, '
                                      'but was passed as a keyword'. \
                                      format(arg=param.name))
      
                  arguments[param_name] = arg_val
      
          if kwargs:
              if kwargs_param is not None:
                  # Process our '**kwargs'-like parameter
                  arguments[kwargs_param.name] = kwargs
              else:
  >               raise TypeError(
                      'got an unexpected keyword argument {arg!r}'.format(
                          arg=next(iter(kwargs))))
  E               TypeError: got an unexpected keyword argument 'extra_kwarg'
  
  /usr/lib/python3.12/inspect.py:3231: TypeError
  ________ TestAttentionValueScaleForward.test_scale_various_values[0.5] _________