"""simple docstring"""



import torch
from torch import nn






class  snake_case__      (       nn.Module       ):


 def __init__( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase=1    ,		lowerCamelCase=False						):
        super().__init__()

        __a       =    n_token
        __a       =    d_embed
        __a       =    d_proj

        __a       =    cutoffs + [n_token]
        __a       =    [0] + self.cutoffs
        __a       =    div_val

        __a       =    self.cutoffs[0]
        __a       =    len(self.cutoffs						) - 1
        __a       =    self.shortlist_size + self.n_clusters

        if self.n_clusters > 0:
               __a       =    nn.Parameter(torch.zeros(self.n_clusters    ,		self.d_embed						)						)
               __a       =    nn.Parameter(torch.zeros(self.n_clusters						)						)

        __a       =    nn.ModuleList()
        __a       =    nn.ParameterList()

        if div_val == 1:
               for i in range(len(self.cutoffs						)						):
                      if d_proj != d_embed:
                             self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCamelCase    ,		lowerCamelCase						)						)						)
                      else:
                             self.out_projs.append(lowerCamelCase						)

               self.out_layers.append(nn.Linear(lowerCamelCase    ,		lowerCamelCase						)						)
        else:
               for i in range(len(self.cutoffs						)						):
                      __a ,   __a       =    self.cutoff_ends[i], self.cutoff_ends[i + 1]
                      __a       =    d_embed // (div_val**i)

                      self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCamelCase    ,		lowerCamelCase						)						)						)

                      self.out_layers.append(nn.Linear(lowerCamelCase    ,		r_idx - l_idx						)						)

        __a       =    keep_order


 def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						):
        if proj is None:
               __a       =    nn.functional.linear(lowerCamelCase    ,		lowerCamelCase    ,		bias=lowerCamelCase						)
        else:
               # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
               __a       =    nn.functional.linear(lowerCamelCase    ,		proj.t().contiguous()						)
               __a       =    nn.functional.linear(lowerCamelCase    ,		lowerCamelCase    ,		bias=lowerCamelCase						)
               # else:
               #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
               #     if bias is not None:
               #         logit = logit + bias

        return logit


 def 						a__   ( self    ,		lowerCamelCase    ,		lowerCamelCase=None    ,		lowerCamelCase=False						):

        if labels is not None:
               # Shift so that tokens < n predict n
               __a       =    hidden[..., :-1, :].contiguous()
               __a       =    labels[..., 1:].contiguous()
               __a       =    hidden.view(-1    ,		hidden.size(-1						)						)
               __a       =    labels.view(-1						)
               if hidden.size(0						) != labels.size(0						):
                      raise RuntimeError("Input and labels should have the same size in the batch dimension."						)
        else:
               __a       =    hidden.view(-1    ,		hidden.size(-1						)						)

        if self.n_clusters == 0:
               __a       =    self._compute_logit(lowerCamelCase    ,		self.out_layers[0].weight    ,		self.out_layers[0].bias    ,		self.out_projs[0]						)
               if labels is not None:
                      __a       =    labels != -100
                      __a       =    torch.zeros_like(lowerCamelCase    ,		dtype=hidden.dtype    ,		device=hidden.device						)
                      __a       =    (
                          -nn.functional.log_softmax(lowerCamelCase    ,		dim=-1						)[mask].gather(1    ,		labels[mask].unsqueeze(1						)						).squeeze(1						)
                      )
               else:
                      __a       =    nn.functional.log_softmax(lowerCamelCase    ,		dim=-1						)
        else:
               # construct weights and biases
               __a ,   __a       =    [], []
               for i in range(len(self.cutoffs						)						):
                      if self.div_val == 1:
                             __a ,   __a       =    self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             __a       =    self.out_layers[0].weight[l_idx:r_idx]
                             __a       =    self.out_layers[0].bias[l_idx:r_idx]
                      else:
                             __a       =    self.out_layers[i].weight
                             __a       =    self.out_layers[i].bias

                      if i == 0:
                             __a       =    torch.cat([weight_i, self.cluster_weight]    ,		dim=0						)
                             __a       =    torch.cat([bias_i, self.cluster_bias]    ,		dim=0						)

                      weights.append(lowerCamelCase						)
                      biases.append(lowerCamelCase						)

               __a ,   __a ,   __a       =    weights[0], biases[0], self.out_projs[0]

               __a       =    self._compute_logit(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						)
               __a       =    nn.functional.log_softmax(lowerCamelCase    ,		dim=1						)

               if labels is None:
                      __a       =    hidden.new_empty((head_logit.size(0						), self.n_token)						)
               else:
                      __a       =    torch.zeros_like(lowerCamelCase    ,		dtype=hidden.dtype    ,		device=hidden.device						)

               __a       =    0
               __a       =    [0] + self.cutoffs
               for i in range(len(lowerCamelCase						) - 1						):
                      __a ,   __a       =    cutoff_values[i], cutoff_values[i + 1]

                      if labels is not None:
                             __a       =    (labels >= l_idx) & (labels < r_idx)
                             __a       =    mask_i.nonzero().squeeze()

                             if indices_i.numel() == 0:
                                    continue

                             __a       =    labels.index_select(0    ,		lowerCamelCase						) - l_idx
                             __a       =    head_logprob.index_select(0    ,		lowerCamelCase						)
                             __a       =    hidden.index_select(0    ,		lowerCamelCase						)
                      else:
                             __a       =    hidden

                      if i == 0:
                             if labels is not None:
                                    __a       =    head_logprob_i.gather(1    ,		target_i[:, None]						).squeeze(1						)
                             else:
                                    __a       =    head_logprob[:, : self.cutoffs[0]]
                      else:
                             __a ,   __a ,   __a       =    weights[i], biases[i], self.out_projs[i]

                             __a       =    self._compute_logit(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						)
                             __a       =    nn.functional.log_softmax(lowerCamelCase    ,		dim=1						)
                             __a       =    self.cutoffs[0] + i - 1  # No probability for the head cluster
                             if labels is not None:
                                    __a       =    head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                                        1    ,		target_i[:, None]						).squeeze(1						)
                             else:
                                    __a       =    head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                                    __a       =    logprob_i

                      if labels is not None:
                             if (hasattr(self    ,		"keep_order"						) and self.keep_order) or keep_order:
                                    out.index_copy_(0    ,		lowerCamelCase    ,		-logprob_i						)
                             else:
                                    out[offset : offset + logprob_i.size(0						)].copy_(-logprob_i						)
                             offset += logprob_i.size(0						)

        return out






 def 						a__   ( self    ,		lowerCamelCase						):
        if self.n_clusters == 0:
               __a       =    self._compute_logit(lowerCamelCase    ,		self.out_layers[0].weight    ,		self.out_layers[0].bias    ,		self.out_projs[0]						)
               return nn.functional.log_softmax(lowerCamelCase    ,		dim=-1						)
        else:
               # construct weights and biases
               __a ,   __a       =    [], []
               for i in range(len(self.cutoffs						)						):
                      if self.div_val == 1:
                             __a ,   __a       =    self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             __a       =    self.out_layers[0].weight[l_idx:r_idx]
                             __a       =    self.out_layers[0].bias[l_idx:r_idx]
                      else:
                             __a       =    self.out_layers[i].weight
                             __a       =    self.out_layers[i].bias

                      if i == 0:
                             __a       =    torch.cat([weight_i, self.cluster_weight]    ,		dim=0						)
                             __a       =    torch.cat([bias_i, self.cluster_bias]    ,		dim=0						)

                      weights.append(lowerCamelCase						)
                      biases.append(lowerCamelCase						)

               __a ,   __a ,   __a       =    weights[0], biases[0], self.out_projs[0]
               __a       =    self._compute_logit(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						)

               __a       =    hidden.new_empty((head_logit.size(0						), self.n_token)						)
               __a       =    nn.functional.log_softmax(lowerCamelCase    ,		dim=1						)

               __a       =    [0] + self.cutoffs
               for i in range(len(lowerCamelCase						) - 1						):
                      __a ,   __a       =    cutoff_values[i], cutoff_values[i + 1]

                      if i == 0:
                             __a       =    head_logprob[:, : self.cutoffs[0]]
                      else:
                             __a ,   __a ,   __a       =    weights[i], biases[i], self.out_projs[i]

                             __a       =    self._compute_logit(lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase    ,		lowerCamelCase						)
                             __a       =    nn.functional.log_softmax(lowerCamelCase    ,		dim=1						)

                             __a       =    head_logprob[:, -i] + tail_logprob_i
                             __a       =    logprob_i

               return out

"""simple docstring"""



import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class  snake_case__      (       unittest.TestCase       ):


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
            "unet/diffusion_pytorch_model.bin",
            # Removed: 'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertFalse(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            # Removed: 'text_encoder/model.safetensors',
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertFalse(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        # pass variant but use the non-variant filenames
        __a       =    [
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
            "unet/diffusion_pytorch_model.fp16.bin",
            # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        __a       =    "fp16"
        self.assertFalse(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        # pass variant but use the non-variant filenames
        __a       =    [
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)






 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            # 'text_encoder/model.fp16.safetensors',
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertFalse(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)
