code: 
  
  
  import inspect
  
  from diffusers.utils import is_flax_available
  from diffusers.utils.testing_utils import require_flax
  
  
  if is_flax_available():
    import jax
  
  
  
  
  
  
  
  @require_flax
  class 					__lowercase					:
  
  
  
    """simple docstring"""
  
    def 				__A      (   self  )     ->   Optional[int]:
  
  
  
  
  
  
  
       '''simple docstring'''
  
  
  
  
  
  
  
       lowerCamelCase		,		lowerCamelCase								=     self.prepare_init_args_and_inputs_for_common()
  
       lowerCamelCase								=     self.model_class(**_lowerCAmelCase  )
       lowerCamelCase								=     model.init(inputs_dict["""prng_key"""]		,			inputs_dict["""sample"""]  )
       jax.lax.stop_gradient(_lowerCAmelCase  )
  
       lowerCamelCase								=     model.apply(_lowerCAmelCase		,			inputs_dict["""sample"""]  )
  
       if isinstance(_lowerCAmelCase		,			_lowerCAmelCase  ):
          lowerCamelCase								=     output.sample
  
       self.assertIsNotNone(_lowerCAmelCase  )
       lowerCamelCase								=     inputs_dict["""sample"""].shape
       self.assertEqual(output.shape		,			_lowerCAmelCase		,			"""Input and output shapes do not match"""  )
  
  
  
  
  
  
  
    def 				__A      (   self  )     ->   List[str]:
  
  
  
  
  
  
  
       '''simple docstring'''
  
  
  
  
  
  
  
       lowerCamelCase		,		lowerCamelCase								=     self.prepare_init_args_and_inputs_for_common()
  
       lowerCamelCase								=     16
       lowerCamelCase								=     (16, 32)
  
       lowerCamelCase								=     self.model_class(**_lowerCAmelCase  )
       lowerCamelCase								=     model.init(inputs_dict["""prng_key"""]		,			inputs_dict["""sample"""]  )
       jax.lax.stop_gradient(_lowerCAmelCase  )
  
       lowerCamelCase								=     model.apply(_lowerCAmelCase		,			inputs_dict["""sample"""]  )
  
       if isinstance(_lowerCAmelCase		,			_lowerCAmelCase  ):
          lowerCamelCase								=     output.sample
  
       self.assertIsNotNone(_lowerCAmelCase  )
       lowerCamelCase								=     inputs_dict["""sample"""].shape
       self.assertEqual(output.shape		,			_lowerCAmelCase		,			"""Input and output shapes do not match"""  )
  
  
  
  
  
  
  
    def 				__A      (   self  )     ->   Any:
  
  
  
  
  
  
  
       '''simple docstring'''
  
  
  
  
  
  
  
       lowerCamelCase								=     """kwargs""" in inspect.signature(self.model_class.__init__  ).parameters
       lowerCamelCase								=     len(self.model_class._deprecated_kwargs  ) > 0
  
       if has_kwarg_in_model_class and not has_deprecated_kwarg:
          raise ValueError(
              F'{self.model_class} has `**kwargs` in its __init__ method but has not defined any deprecated kwargs'
              """ under the `_deprecated_kwargs` class attribute. Make sure to either remove `**kwargs` if there are"""
              """ no deprecated arguments or add the deprecated argument with `_deprecated_kwargs ="""
              """ [<deprecated_argument>]`"""  )
  
       if not has_kwarg_in_model_class and has_deprecated_kwarg:
          raise ValueError(
              F'{self.model_class} doesn\'t have `**kwargs` in its __init__ method but has defined deprecated kwargs'
              """ under the `_deprecated_kwargs` class attribute. Make sure to either add the `**kwargs` argument to"""
              F' {self.model_class}.__init__ if there are deprecated arguments or remove the deprecated argument'
              """ from `_deprecated_kwargs = [<deprecated_argument>]`"""  )
  
  
style_context: 
  
  
  def   __lowerCamelCase	(       lowerCamelCase__		:	int = 100     ,       lowerCamelCase__		:	int = 100			):
  
  
  
  
  
  			'''simple docstring'''
  
  
  
  
  			return max(
  			    sum(int(lowerCamelCase__			) for x in str(base**power			)			)
  			    for base in range(lowerCamelCase__			)
  			    for power in range(lowerCamelCase__			)			)
  
  
  # Tests
  if __name__ == "__main__":
  		import doctest
  
  		doctest.testmod()
  
  