



import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule






class SCREAMING_SNAKE_CASE__	(						SCREAMING_SNAKE_CASE__       ):




    '''simple docstring'''






    __lowerCamelCase     :					Tuple   = 42
    __lowerCamelCase     :					Dict   = jnp.floataa
    __lowerCamelCase     :					str   = True






    def      _lowerCAmelCase   (			self						):
        super().setup()
        A					:			List[Any]		=  nn.Dense(5,						dtype=self.dtype						)






    def __call__(			self,						*lowerCamelCase__,						**lowerCamelCase__						):
        A					:			Optional[Any]		=  super().__call__(*snake_case_,						**snake_case_						)
        A					:			Tuple		=  self.cls(outputs[2]						)
        return outputs[:2] + (cls_out,)






class SCREAMING_SNAKE_CASE__	(						SCREAMING_SNAKE_CASE__       ):




    '''simple docstring'''






    __lowerCamelCase     :					List[str]   = FlaxBigBirdForNaturalQuestionsModule




def       __UpperCamelCase      (    _lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase						)			->    Dict:



    """simple docstring"""






    def cross_entropy(_lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase=None						):
        A					:			Tuple		=  logits.shape[-1]
        A					:			Optional[Any]		=  (labels[..., None] == jnp.arange(__lowerCamelCase						)[None]).astype("""f4"""						)
        A					:			Optional[Any]		=  jax.nn.log_softmax(__lowerCamelCase ,	axis=-1						)
        A					:			Union[str, Any]		=  -jnp.sum(labels * logits ,	axis=-1						)
        if reduction is not None:
            A					:			List[str]		=  reduction(__lowerCamelCase						)
        return loss

    A					:			Optional[int]		=  partial(__lowerCamelCase ,	reduction=jnp.mean						)
    A					:			List[Any]		=  cross_entropy(__lowerCamelCase ,	__lowerCamelCase						)
    A					:			int		=  cross_entropy(__lowerCamelCase ,	__lowerCamelCase						)
    A					:			List[Any]		=  cross_entropy(__lowerCamelCase ,	__lowerCamelCase						)
    return (start_loss + end_loss + pooled_loss) / 3






@dataclass
class SCREAMING_SNAKE_CASE__	:




    '''simple docstring'''






    __lowerCamelCase     :					Any   = "google/bigbird-roberta-base"
    __lowerCamelCase     :					Union[str, Any]   = 3000
    __lowerCamelCase     :					List[Any]   = 1_0500

    __lowerCamelCase     :					List[Any]   = 128
    __lowerCamelCase     :					Optional[int]   = 3

    __lowerCamelCase     :					Dict   = 1
    __lowerCamelCase     :					Any   = 5

    # tx_args
    __lowerCamelCase     :					List[Any]   = 3e-5
    __lowerCamelCase     :					Any   = 0.0
    __lowerCamelCase     :					Tuple   = 2_0000
    __lowerCamelCase     :					str   = 0.0_095

    __lowerCamelCase     :					Dict   = "bigbird-roberta-natural-questions"
    __lowerCamelCase     :					int   = "training-expt"
    __lowerCamelCase     :					str   = "data/nq-training.jsonl"
    __lowerCamelCase     :					int   = "data/nq-validation.jsonl"






    def      _lowerCAmelCase   (			self						):
        os.makedirs(self.base_dir,						exist_ok=snake_case_						)
        A					:			List[str]		=  os.path.join(self.base_dir,						self.save_dir						)
        A					:			Optional[int]		=  self.batch_size_per_device * jax.device_count()






@dataclass
class SCREAMING_SNAKE_CASE__	:




    '''simple docstring'''






    __lowerCamelCase     :					Tuple   = 42
    __lowerCamelCase     :					Dict   = 4096  # no dynamic padding on TPUs






    def __call__(			self,						lowerCamelCase__						):
        A					:			int		=  self.collate_fn(snake_case_						)
        A					:			int		=  jax.tree_util.tree_map(snake_case_,						snake_case_						)
        return batch






    def      _lowerCAmelCase   (			self,						lowerCamelCase__						):
        A	,				A					:			Any		=  self.fetch_inputs(features["""input_ids"""]						)
        A					:			int		=  {
            """input_ids""": jnp.array(snake_case_,						dtype=jnp.intaa						),
            """attention_mask""": jnp.array(snake_case_,						dtype=jnp.intaa						),
            """start_labels""": jnp.array(features["""start_token"""],						dtype=jnp.intaa						),
            """end_labels""": jnp.array(features["""end_token"""],						dtype=jnp.intaa						),
            """pooled_labels""": jnp.array(features["""category"""],						dtype=jnp.intaa						),
        }
        return batch






    def      _lowerCAmelCase   (			self,						lowerCamelCase__						):
        A					:			Optional[int]		=  [self._fetch_inputs(snake_case_						) for ids in input_ids]
        return zip(*snake_case_						)






    def      _lowerCAmelCase   (			self,						lowerCamelCase__						):
        A					:			Union[str, Any]		=  [1 for _ in range(len(snake_case_						)						)]
        while len(snake_case_						) < self.max_length:
            input_ids.append(self.pad_id						)
            attention_mask.append(0						)
        return input_ids, attention_mask




def       __UpperCamelCase      (    _lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase=None						)			->    Optional[Any]:



    """simple docstring"""






    if seed is not None:
        A					:			Optional[Any]		=  dataset.shuffle(seed=__lowerCamelCase						)
    for i in range(len(__lowerCamelCase						) // batch_size						):
        A					:			Any		=  dataset[i * batch_size : (i + 1) * batch_size]
        yield dict(__lowerCamelCase						)




@partial(jax.pmap ,	axis_name="""batch"""						)
def       __UpperCamelCase      (    _lowerCAmelCase ,	_lowerCAmelCase ,	**_lowerCAmelCase						)			->    Union[str, Any]:



    """simple docstring"""






    def loss_fn(_lowerCAmelCase						):
        A					:			List[Any]		=  model_inputs.pop("""start_labels"""						)
        A					:			Any		=  model_inputs.pop("""end_labels"""						)
        A					:			List[str]		=  model_inputs.pop("""pooled_labels"""						)

        A					:			Optional[int]		=  state.apply_fn(**__lowerCamelCase ,	params=__lowerCamelCase ,	dropout_rng=__lowerCamelCase ,	train=__lowerCamelCase						)
        A	,				A	,				A					:			List[str]		=  outputs

        return state.loss_fn(
            __lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	)

    A	,				A					:			Optional[int]		=  jax.random.split(__lowerCamelCase						)
    A					:			Optional[Any]		=  jax.value_and_grad(__lowerCamelCase						)
    A	,				A					:			str		=  grad_fn(state.params						)
    A					:			List[str]		=  jax.lax.pmean({"""loss""": loss} ,	axis_name="""batch"""						)
    A					:			int		=  jax.lax.pmean(__lowerCamelCase ,	"""batch"""						)

    A					:			str		=  state.apply_gradients(grads=__lowerCamelCase						)
    return state, metrics, new_drp_rng




@partial(jax.pmap ,	axis_name="""batch"""						)
def       __UpperCamelCase      (    _lowerCAmelCase ,	**_lowerCAmelCase						)			->    List[str]:



    """simple docstring"""






    A					:			Union[str, Any]		=  model_inputs.pop("""start_labels"""						)
    A					:			int		=  model_inputs.pop("""end_labels"""						)
    A					:			Optional[Any]		=  model_inputs.pop("""pooled_labels"""						)

    A					:			Dict		=  state.apply_fn(**__lowerCamelCase ,	params=state.params ,	train=__lowerCamelCase						)
    A	,				A	,				A					:			Optional[int]		=  outputs

    A					:			Optional[Any]		=  state.loss_fn(__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase						)
    A					:			Tuple		=  jax.lax.pmean({"""loss""": loss} ,	axis_name="""batch"""						)
    return metrics






class SCREAMING_SNAKE_CASE__	(						train_state.TrainState       ):




    '''simple docstring'''






    __lowerCamelCase     :					str   = struct.field(pytree_node=SCREAMING_SNAKE_CASE__       )






@dataclass
class SCREAMING_SNAKE_CASE__	:




    '''simple docstring'''






    __lowerCamelCase     :					Optional[int]   = 42
    __lowerCamelCase     :					Optional[Any]   = 42
    __lowerCamelCase     :					Any   = 42
    __lowerCamelCase     :					List[Any]   = 42
    __lowerCamelCase     :					List[str]   = 42
    __lowerCamelCase     :					Union[str, Any]   = 42
    __lowerCamelCase     :					List[Any]   = None






    def      _lowerCAmelCase   (			self,						lowerCamelCase__,						lowerCamelCase__,						lowerCamelCase__,						lowerCamelCase__=None						):
        A					:			Union[str, Any]		=  model.params
        A					:			Tuple		=  TrainState.create(
            apply_fn=model.__call__,						params=snake_case_,						tx=snake_case_,						loss_fn=snake_case_,						)
        if ckpt_dir is not None:
            A	,				A	,				A	,				A	,				A					:			List[str]		=  restore_checkpoint(snake_case_,						snake_case_						)
            A					:			List[Any]		=  {
                """lr""": args.lr,
                """init_lr""": args.init_lr,
                """warmup_steps""": args.warmup_steps,
                """num_train_steps""": num_train_steps,
                """weight_decay""": args.weight_decay,
            }
            A	,				A					:			List[str]		=  build_tx(**snake_case_						)
            A					:			str		=  train_state.TrainState(
                step=snake_case_,						apply_fn=model.__call__,						params=snake_case_,						tx=snake_case_,						opt_state=snake_case_,						)
            A					:			Optional[int]		=  args
            A					:			List[Any]		=  data_collator
            A					:			Optional[Any]		=  lr
            A					:			List[str]		=  params
        A					:			Optional[Any]		=  jax_utils.replicate(snake_case_						)
        return state






    def      _lowerCAmelCase   (			self,						lowerCamelCase__,						lowerCamelCase__,						lowerCamelCase__						):
        A					:			Optional[int]		=  self.args
        A					:			Optional[int]		=  len(snake_case_						) // args.batch_size

        A					:			int		=  jax.random.PRNGKey(0						)
        A					:			Optional[Any]		=  jax.random.split(snake_case_,						jax.device_count()						)
        for epoch in range(args.max_epochs						):
            A					:			Dict		=  jnp.array(0,						dtype=jnp.floataa						)
            A					:			List[str]		=  get_batched_dataset(snake_case_,						args.batch_size,						seed=snake_case_						)
            A					:			Any		=  0
            for batch in tqdm(snake_case_,						total=snake_case_,						desc=f'''Running EPOCH-{epoch}'''						):
                A					:			str		=  self.data_collator(snake_case_						)
                A	,				A	,				A					:			List[str]		=  self.train_step_fn(snake_case_,						snake_case_,						**snake_case_						)
                running_loss += jax_utils.unreplicate(metrics["""loss"""]						)
                i += 1
                if i % args.logging_steps == 0:
                    A					:			List[str]		=  jax_utils.unreplicate(state.step						)
                    A					:			Union[str, Any]		=  running_loss.item() / i
                    A					:			Tuple		=  self.scheduler_fn(state_step - 1						)

                    A					:			str		=  self.evaluate(snake_case_,						snake_case_						)
                    A					:			Any		=  {
                        """step""": state_step.item(),
                        """eval_loss""": eval_loss.item(),
                        """tr_loss""": tr_loss,
                        """lr""": lr.item(),
                    }
                    tqdm.write(str(snake_case_						)						)
                    self.logger.log(snake_case_,						commit=snake_case_						)

                if i % args.save_steps == 0:
                    self.save_checkpoint(args.save_dir + f'''-e{epoch}-s{i}''',						state=snake_case_						)






    def      _lowerCAmelCase   (			self,						lowerCamelCase__,						lowerCamelCase__						):
        A					:			Any		=  get_batched_dataset(snake_case_,						self.args.batch_size						)
        A					:			int		=  len(snake_case_						) // self.args.batch_size
        A					:			Optional[int]		=  jnp.array(0,						dtype=jnp.floataa						)
        A					:			Optional[int]		=  0
        for batch in tqdm(snake_case_,						total=snake_case_,						desc="""Evaluating ... """						):
            A					:			Dict		=  self.data_collator(snake_case_						)
            A					:			Optional[int]		=  self.val_step_fn(snake_case_,						**snake_case_						)
            running_loss += jax_utils.unreplicate(metrics["""loss"""]						)
            i += 1
        return running_loss / i






    def      _lowerCAmelCase   (			self,						lowerCamelCase__,						lowerCamelCase__						):
        A					:			List[str]		=  jax_utils.unreplicate(snake_case_						)
        print(f'''SAVING CHECKPOINT IN {save_dir}''',						end=""" ... """						)
        self.model_save_fn(snake_case_,						params=state.params						)
        with open(os.path.join(snake_case_,						"""opt_state.msgpack"""						),						"""wb"""						) as f:
            f.write(to_bytes(state.opt_state						)						)
        joblib.dump(self.args,						os.path.join(snake_case_,						"""args.joblib"""						)						)
        joblib.dump(self.data_collator,						os.path.join(snake_case_,						"""data_collator.joblib"""						)						)
        with open(os.path.join(snake_case_,						"""training_state.json"""						),						"""w"""						) as f:
            json.dump({"""step""": state.step.item()},						snake_case_						)
        print("""DONE"""						)




def       __UpperCamelCase      (    _lowerCAmelCase ,	_lowerCAmelCase						)			->    Tuple:



    """simple docstring"""






    print(f'''RESTORING CHECKPOINT FROM {save_dir}''' ,	end=""" ... """						)
    with open(os.path.join(__lowerCamelCase ,	"""flax_model.msgpack"""						) ,	"""rb"""						) as f:
        A					:			Optional[int]		=  from_bytes(state.params ,	f.read()						)

    with open(os.path.join(__lowerCamelCase ,	"""opt_state.msgpack"""						) ,	"""rb"""						) as f:
        A					:			Tuple		=  from_bytes(state.opt_state ,	f.read()						)

    A					:			List[str]		=  joblib.load(os.path.join(__lowerCamelCase ,	"""args.joblib"""						)						)
    A					:			Union[str, Any]		=  joblib.load(os.path.join(__lowerCamelCase ,	"""data_collator.joblib"""						)						)

    with open(os.path.join(__lowerCamelCase ,	"""training_state.json"""						) ,	"""r"""						) as f:
        A					:			Union[str, Any]		=  json.load(__lowerCamelCase						)
    A					:			Any		=  training_state["""step"""]

    print("""DONE"""						)
    return params, opt_state, step, args, data_collator




def       __UpperCamelCase      (    _lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase						)			->    Optional[Any]:



    """simple docstring"""






    A					:			Tuple		=  num_train_steps - warmup_steps
    A					:			List[Any]		=  optax.linear_schedule(init_value=__lowerCamelCase ,	end_value=__lowerCamelCase ,	transition_steps=__lowerCamelCase						)
    A					:			Any		=  optax.linear_schedule(init_value=__lowerCamelCase ,	end_value=1e-7 ,	transition_steps=__lowerCamelCase						)
    A					:			Optional[Any]		=  optax.join_schedules(schedules=[warmup_fn, decay_fn] ,	boundaries=[warmup_steps]						)
    return lr




def       __UpperCamelCase      (    _lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase ,	_lowerCAmelCase						)			->    int:



    """simple docstring"""






    def weight_decay_mask(_lowerCAmelCase						):
        A					:			List[str]		=  traverse_util.flatten_dict(__lowerCamelCase						)
        A					:			Tuple		=  {k: (v[-1] != """bias""" and v[-2:] != ("""LayerNorm""", """scale""")) for k, v in params.items()}
        return traverse_util.unflatten_dict(__lowerCamelCase						)

    A					:			Dict		=  scheduler_fn(__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase ,	__lowerCamelCase						)

    A					:			str		=  optax.adamw(learning_rate=__lowerCamelCase ,	weight_decay=__lowerCamelCase ,	mask=__lowerCamelCase						)
    return tx, lr







import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
				import torch

				from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin






class     __a			:



	def __init__(    self				:     Union[str, Any]					,						snake_case_				:     Dict					,						snake_case_				:     Tuple=None					,						snake_case_				:     List[str]=None					,						snake_case_				:     Union[str, Any]=None					,						snake_case_				:     str="resnet50"					,						snake_case_				:     List[Any]=3					,						snake_case_				:     Optional[int]=32					,						snake_case_				:     Union[str, Any]=3					,						snake_case_				:     Tuple=True					,						snake_case_				:     List[str]=True					,						)->				Optional[Any]:
						__lowerCAmelCase						 =parent
						__lowerCAmelCase						 =out_indices if out_indices is not None else [4]
						__lowerCAmelCase						 =stage_names
						__lowerCAmelCase						 =out_features
						__lowerCAmelCase						 =backbone
						__lowerCAmelCase						 =batch_size
						__lowerCAmelCase						 =image_size
						__lowerCAmelCase						 =num_channels
						__lowerCAmelCase						 =use_pretrained_backbone
						__lowerCAmelCase						 =is_training



	def 		UpperCamelCase				(    self				:     int)->				Dict:
						__lowerCAmelCase						 =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])
						__lowerCAmelCase						 =self.get_config()

						return config, pixel_values



	def 		UpperCamelCase				(    self				:     Optional[int])->				Optional[int]:
						return TimmBackboneConfig(
						    image_size=self.image_size					,						num_channels=self.num_channels					,						out_features=self.out_features					,						out_indices=self.out_indices					,						stage_names=self.stage_names					,						use_pretrained_backbone=self.use_pretrained_backbone					,						backbone=self.backbone					,						)



	def 		UpperCamelCase				(    self				:     str					,						snake_case_				:     int					,						snake_case_				:     Union[str, Any])->				str:
						__lowerCAmelCase						 =TimmBackbone(config=snake_case_)
						model.to(snake_case_)
						model.eval()
						with torch.no_grad():
											__lowerCAmelCase						 =model(snake_case_)
						self.parent.assertEqual(
						    result.feature_map[-1].shape					,						(self.batch_size, model.channels[-1], 14, 14)					,						)



	def 		UpperCamelCase				(    self				:     List[str])->				Union[str, Any]:
						__lowerCAmelCase						 =self.prepare_config_and_inputs()
						__lowerCAmelCase ,						__lowerCAmelCase						 =config_and_inputs
						__lowerCAmelCase						 ={"""pixel_values""": pixel_values}
						return config, inputs_dict






@require_torch
@require_timm
class     __a			(  SCREAMING_SNAKE_CASE						,     SCREAMING_SNAKE_CASE						,     SCREAMING_SNAKE_CASE						,     unittest.TestCase					):
	SCREAMING_SNAKE_CASE  	=							(TimmBackbone,) if is_torch_available() else ()
	SCREAMING_SNAKE_CASE  	=							{"feature-extraction": TimmBackbone} if is_torch_available() else {}
	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False



	def 		UpperCamelCase				(    self				:     Union[str, Any])->				str:
						__lowerCAmelCase						 =TimmBackboneModelTester(self)
						__lowerCAmelCase						 =ConfigTester(self					,						config_class=snake_case_					,						has_text_modality=snake_case_)



	def 		UpperCamelCase				(    self				:     Tuple)->				Optional[Any]:
						self.config_tester.create_and_test_config_to_json_string()
						self.config_tester.create_and_test_config_to_json_file()
						self.config_tester.create_and_test_config_from_and_save_pretrained()
						self.config_tester.create_and_test_config_with_num_labels()
						self.config_tester.check_config_can_be_init_without_params()
						self.config_tester.check_config_arguments_init()



	def 		UpperCamelCase				(    self				:     Any)->				Dict:
						__lowerCAmelCase						 ="""resnet18"""
						__lowerCAmelCase						 ="""microsoft/resnet-18"""

						__lowerCAmelCase						 =AutoBackbone.from_pretrained(snake_case_					,						use_timm_backbone=snake_case_)
						__lowerCAmelCase						 =AutoBackbone.from_pretrained(snake_case_)

						self.assertEqual(len(timm_model.out_features)					,						len(transformers_model.out_features))
						self.assertEqual(len(timm_model.stage_names)					,						len(transformers_model.stage_names))
						self.assertEqual(timm_model.channels					,						transformers_model.channels)
						# Out indices are set to the last layer by default. For timm models, we don't know
						# the number of layers in advance, so we set it to (-1,), whereas for transformers
						# models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
						self.assertEqual(timm_model.out_indices					,						(-1,))
						self.assertEqual(transformers_model.out_indices					,						[len(timm_model.stage_names) - 1])

						__lowerCAmelCase						 =AutoBackbone.from_pretrained(snake_case_					,						use_timm_backbone=snake_case_					,						out_indices=[1, 2, 3])
						__lowerCAmelCase						 =AutoBackbone.from_pretrained(snake_case_					,						out_indices=[1, 2, 3])

						self.assertEqual(timm_model.out_indices					,						transformers_model.out_indices)
						self.assertEqual(len(timm_model.out_features)					,						len(transformers_model.out_features))
						self.assertEqual(timm_model.channels					,						transformers_model.channels)



	@unittest.skip("""TimmBackbone doesn't support feed forward chunking""")
	def 		UpperCamelCase				(    self				:     Dict)->				Any:
						pass



	@unittest.skip("""TimmBackbone doesn't have num_hidden_layers attribute""")
	def 		UpperCamelCase				(    self				:     Tuple)->				Dict:
						pass



	@unittest.skip("""TimmBackbone initialization is managed on the timm side""")
	def 		UpperCamelCase				(    self				:     Union[str, Any])->				List[str]:
						pass



	@unittest.skip("""TimmBackbone models doesn't have inputs_embeds""")
	def 		UpperCamelCase				(    self				:     List[str])->				List[str]:
						pass



	@unittest.skip("""TimmBackbone models doesn't have inputs_embeds""")
	def 		UpperCamelCase				(    self				:     int)->				int:
						pass



	@unittest.skip("""TimmBackbone model cannot be created without specifying a backbone checkpoint""")
	def 		UpperCamelCase				(    self				:     Dict)->				int:
						pass



	@unittest.skip("""Only checkpoints on timm can be loaded into TimmBackbone""")
	def 		UpperCamelCase				(    self				:     Any)->				Dict:
						pass



	@unittest.skip("""model weights aren't tied in TimmBackbone.""")
	def 		UpperCamelCase				(    self				:     int)->				Tuple:
						pass



	@unittest.skip("""model weights aren't tied in TimmBackbone.""")
	def 		UpperCamelCase				(    self				:     Optional[int])->				Union[str, Any]:
						pass



	@unittest.skip("""Only checkpoints on timm can be loaded into TimmBackbone""")
	def 		UpperCamelCase				(    self				:     Optional[Any])->				Optional[Any]:
						pass



	@unittest.skip("""Only checkpoints on timm can be loaded into TimmBackbone""")
	def 		UpperCamelCase				(    self				:     Any)->				List[Any]:
						pass



	@unittest.skip("""TimmBackbone doesn't have hidden size info in its configuration.""")
	def 		UpperCamelCase				(    self				:     Any)->				Tuple:
						pass



	@unittest.skip("""TimmBackbone doesn't support output_attentions.""")
	def 		UpperCamelCase				(    self				:     List[Any])->				str:
						pass



	@unittest.skip("""Safetensors is not supported by timm.""")
	def 		UpperCamelCase				(    self				:     Dict)->				Optional[int]:
						pass



	@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""")
	def 		UpperCamelCase				(    self				:     Optional[int])->				Tuple:
						pass



	def 		UpperCamelCase				(    self				:     int)->				Union[str, Any]:
						__lowerCAmelCase ,						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
											__lowerCAmelCase						 =model_class(snake_case_)
											__lowerCAmelCase						 =inspect.signature(model.forward)
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											__lowerCAmelCase						 =[*signature.parameters.keys()]

											__lowerCAmelCase						 =["""pixel_values"""]
											self.assertListEqual(arg_names[:1]					,						snake_case_)



	def 		UpperCamelCase				(    self				:     Dict)->				Optional[Any]:
						__lowerCAmelCase ,						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs_for_common()
						__lowerCAmelCase						 =True
						__lowerCAmelCase						 =self.has_attentions

						# no need to test all models as different heads yield the same functionality
						__lowerCAmelCase						 =self.all_model_classes[0]
						__lowerCAmelCase						 =model_class(snake_case_)
						model.to(snake_case_)

						__lowerCAmelCase						 =self._prepare_for_class(snake_case_					,						snake_case_)
						__lowerCAmelCase						 =model(**snake_case_)
						__lowerCAmelCase						 =outputs[0][-1]

						# Encoder-/Decoder-only models
						__lowerCAmelCase						 =outputs.hidden_states[0]
						hidden_states.retain_grad()

						if self.has_attentions:
											__lowerCAmelCase						 =outputs.attentions[0]
											attentions.retain_grad()

						output.flatten()[0].backward(retain_graph=snake_case_)

						self.assertIsNotNone(hidden_states.grad)

						if self.has_attentions:
											self.assertIsNotNone(attentions.grad)



	def 		UpperCamelCase				(    self				:     Tuple)->				List[Any]:
						__lowerCAmelCase ,						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
											__lowerCAmelCase						 =model_class(snake_case_)
											model.to(snake_case_)
											model.eval()
											__lowerCAmelCase						 =model(**snake_case_)

											self.assertEqual(len(result.feature_maps)					,						len(config.out_indices))
											self.assertEqual(len(model.channels)					,						len(config.out_indices))

											# Check output of last stage is taken if out_features=None, out_indices=None
											__lowerCAmelCase						 =copy.deepcopy(snake_case_)
											__lowerCAmelCase						 =None
											__lowerCAmelCase						 =model_class(snake_case_)
											model.to(snake_case_)
											model.eval()
											__lowerCAmelCase						 =model(**snake_case_)

											self.assertEqual(len(result.feature_maps)					,						1)
											self.assertEqual(len(model.channels)					,						1)

											# Check backbone can be initialized with fresh weights
											__lowerCAmelCase						 =copy.deepcopy(snake_case_)
											__lowerCAmelCase						 =False
											__lowerCAmelCase						 =model_class(snake_case_)
											model.to(snake_case_)
											model.eval()
											__lowerCAmelCase						 =model(**snake_case_)




