






import tensorflow as tf

from ...tf_utils import shape_list

class     _SCREAMING_SNAKE_CASE       (			tf.keras.layers.Layer      ):





	def __init__(		self				:					Tuple ,  __lowerCamelCase				:					Optional[int] ,  __lowerCamelCase				:					Dict ,  __lowerCamelCase				:					List[Any] ,  __lowerCamelCase				:					Dict ,  __lowerCamelCase				:					Optional[int]=1 ,  __lowerCamelCase				:					Tuple=False ,  **__lowerCamelCase				:					Dict	):
		super().__init__(**__lowerCamelCase	)

		UpperCamelCase     :str          =	vocab_size
		UpperCamelCase     :Union[str, Any]          =	d_embed
		UpperCamelCase     :Optional[Any]          =	d_proj

		UpperCamelCase     :Union[str, Any]          =	cutoffs + [vocab_size]
		UpperCamelCase     :int          =	[0] + self.cutoffs
		UpperCamelCase     :List[str]          =	div_val

		UpperCamelCase     :Optional[Any]          =	self.cutoffs[0]
		UpperCamelCase     :Optional[int]          =	len(self.cutoffs	) - 1
		UpperCamelCase     :Tuple          =	self.shortlist_size + self.n_clusters
		UpperCamelCase     :List[str]          =	keep_order

		UpperCamelCase     :str          =	[]
		UpperCamelCase     :Optional[Any]          =	[]





	def 						_A (		self				:					Tuple ,  __lowerCamelCase				:					List[str]	):
		if self.n_clusters > 0:
			UpperCamelCase     :List[str]          =	self.add_weight(
			    shape=(self.n_clusters, self.d_embed) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name="""cluster_weight"""	)
			UpperCamelCase     :Union[str, Any]          =	self.add_weight(
			    shape=(self.n_clusters,) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name="""cluster_bias"""	)

		if self.div_val == 1:
			for i in range(len(self.cutoffs	)	):
				if self.d_proj != self.d_embed:
					UpperCamelCase     :Optional[Any]          =	self.add_weight(
					    shape=(self.d_embed, self.d_proj) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name=F"""out_projs_._{i}""" ,  )
					self.out_projs.append(__lowerCamelCase	)
				else:
					self.out_projs.append(__lowerCamelCase	)
				UpperCamelCase     :List[str]          =	self.add_weight(
				    shape=(self.vocab_size, self.d_embed) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name=F"""out_layers_._{i}_._weight""" ,  )
				UpperCamelCase     :Dict          =	self.add_weight(
				    shape=(self.vocab_size,) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name=F"""out_layers_._{i}_._bias""" ,  )
				self.out_layers.append((weight, bias)	)
		else:
			for i in range(len(self.cutoffs	)	):
				UpperCamelCase     ,      UpperCamelCase     :int          =	self.cutoff_ends[i], self.cutoff_ends[i + 1]
				UpperCamelCase     :Optional[int]          =	self.d_embed // (self.div_val**i)

				UpperCamelCase     :Tuple          =	self.add_weight(
				    shape=(d_emb_i, self.d_proj) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name=F"""out_projs_._{i}"""	)
				self.out_projs.append(__lowerCamelCase	)
				UpperCamelCase     :Optional[Any]          =	self.add_weight(
				    shape=(r_idx - l_idx, d_emb_i) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name=F"""out_layers_._{i}_._weight""" ,  )
				UpperCamelCase     :List[str]          =	self.add_weight(
				    shape=(r_idx - l_idx,) ,  initializer="""zeros""" ,  trainable=__lowerCamelCase ,  name=F"""out_layers_._{i}_._bias""" ,  )
				self.out_layers.append((weight, bias)	)
		super().build(__lowerCamelCase	)





	@staticmethod
	def 						_A (		__lowerCamelCase				:					List[Any] ,  __lowerCamelCase				:					Dict ,  __lowerCamelCase				:					Optional[int] ,  __lowerCamelCase				:					Optional[Any]=None	):
		UpperCamelCase     :Tuple          =	x
		if proj is not None:
			UpperCamelCase     :Any          =	tf.einsum("""ibd,ed->ibe""" ,  __lowerCamelCase ,  __lowerCamelCase	)
		return tf.einsum("""ibd,nd->ibn""" ,  __lowerCamelCase ,  __lowerCamelCase	) + b





	@staticmethod
	def 						_A (		__lowerCamelCase				:					List[Any] ,  __lowerCamelCase				:					List[Any]	):
		UpperCamelCase     :Any          =	shape_list(__lowerCamelCase	)
		UpperCamelCase     :Optional[Any]          =	tf.range(lp_size[0] ,  dtype=target.dtype	)
		UpperCamelCase     :str          =	tf.stack([r, target] ,  1	)
		return tf.gather_nd(__lowerCamelCase ,  __lowerCamelCase	)





	def 						_A (		self				:					str ,  __lowerCamelCase				:					Tuple ,  __lowerCamelCase				:					Tuple ,  __lowerCamelCase				:					int=True ,  __lowerCamelCase				:					List[Any]=False	):
		UpperCamelCase     :List[Any]          =	0
		if self.n_clusters == 0:
			UpperCamelCase     :List[str]          =	self._logit(__lowerCamelCase ,  self.out_layers[0][0] ,  self.out_layers[0][1] ,  self.out_projs[0]	)
			if target is not None:
				UpperCamelCase     :Tuple          =	tf.nn.sparse_softmax_cross_entropy_with_logits(labels=__lowerCamelCase ,  logits=__lowerCamelCase	)
			UpperCamelCase     :int          =	tf.nn.log_softmax(__lowerCamelCase ,  axis=-1	)
		else:
			UpperCamelCase     :List[str]          =	shape_list(__lowerCamelCase	)
			UpperCamelCase     :List[Any]          =	[]
			UpperCamelCase     :Union[str, Any]          =	tf.zeros(hidden_sizes[:2]	)
			for i in range(len(self.cutoffs	)	):
				UpperCamelCase     ,      UpperCamelCase     :str          =	self.cutoff_ends[i], self.cutoff_ends[i + 1]
				if target is not None:
					UpperCamelCase     :int          =	(target >= l_idx) & (target < r_idx)
					UpperCamelCase     :int          =	tf.where(__lowerCamelCase	)
					UpperCamelCase     :Tuple          =	tf.boolean_mask(__lowerCamelCase ,  __lowerCamelCase	) - l_idx

				if self.div_val == 1:
					UpperCamelCase     :Tuple          =	self.out_layers[0][0][l_idx:r_idx]
					UpperCamelCase     :Union[str, Any]          =	self.out_layers[0][1][l_idx:r_idx]
				else:
					UpperCamelCase     :List[str]          =	self.out_layers[i][0]
					UpperCamelCase     :Union[str, Any]          =	self.out_layers[i][1]

				if i == 0:
					UpperCamelCase     :str          =	tf.concat([cur_W, self.cluster_weight] ,  0	)
					UpperCamelCase     :List[str]          =	tf.concat([cur_b, self.cluster_bias] ,  0	)

					UpperCamelCase     :Union[str, Any]          =	self._logit(__lowerCamelCase ,  __lowerCamelCase ,  __lowerCamelCase ,  self.out_projs[0]	)
					UpperCamelCase     :List[Any]          =	tf.nn.log_softmax(__lowerCamelCase	)
					out.append(head_logprob[..., : self.cutoffs[0]]	)
					if target is not None:
						UpperCamelCase     :str          =	tf.boolean_mask(__lowerCamelCase ,  __lowerCamelCase	)
						UpperCamelCase     :Tuple          =	self._gather_logprob(__lowerCamelCase ,  __lowerCamelCase	)
				else:
					UpperCamelCase     :int          =	self._logit(__lowerCamelCase ,  __lowerCamelCase ,  __lowerCamelCase ,  self.out_projs[i]	)
					UpperCamelCase     :Optional[int]          =	tf.nn.log_softmax(__lowerCamelCase	)
					UpperCamelCase     :Dict          =	self.cutoffs[0] + i - 1  # No probability for the head cluster
					UpperCamelCase     :List[str]          =	head_logprob[..., cluster_prob_idx, None] + tail_logprob
					out.append(__lowerCamelCase	)
					if target is not None:
						UpperCamelCase     :List[Any]          =	tf.boolean_mask(__lowerCamelCase ,  __lowerCamelCase	)
						UpperCamelCase     :List[str]          =	tf.boolean_mask(__lowerCamelCase ,  __lowerCamelCase	)
						UpperCamelCase     :Any          =	self._gather_logprob(__lowerCamelCase ,  __lowerCamelCase	)
						cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
				if target is not None:
					loss += tf.scatter_nd(__lowerCamelCase ,  -cur_logprob ,  shape_list(__lowerCamelCase	)	)
			UpperCamelCase     :Dict          =	tf.concat(__lowerCamelCase ,  axis=-1	)

		if target is not None:
			if return_mean:
				UpperCamelCase     :Dict          =	tf.reduce_mean(__lowerCamelCase	)
			# Add the training-time loss value to the layer using `self.add_loss()`.
			self.add_loss(__lowerCamelCase	)

			# Log the loss as a metric (we could log arbitrary metrics,
			# including different metrics for training and inference.
			self.add_metric(__lowerCamelCase ,  name=self.name ,  aggregation="""mean""" if return_mean else """"""	)

		return out








from collections.abc import Generator
from math import sin

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			bytes:






	"""simple docstring"""



	if len(__magic_name__					) != 32:
		raise ValueError("""Input must be of length 32"""					)

	UpperCamelCase     :int          =	B""""""
	for i in [3, 2, 1, 0]:
		little_endian += string_aa[8 * i : 8 * i + 8]
	return little_endian

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int					)		->			bytes:






	"""simple docstring"""



	if i < 0:
		raise ValueError("""Input must be non-negative"""					)

	UpperCamelCase     :Any          =	format(__magic_name__	,			"""08x"""					)[-8:]
	UpperCamelCase     :Union[str, Any]          =	B""""""
	for i in [3, 2, 1, 0]:
		little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode("""utf-8"""					)
	return little_endian_hex

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			bytes:






	"""simple docstring"""



	UpperCamelCase     :str          =	B""""""
	for char in message:
		bit_string += format(__magic_name__	,			"""08b"""					).encode("""utf-8"""					)
	UpperCamelCase     :Any          =	format(len(__magic_name__					)	,			"""064b"""					).encode("""utf-8"""					)

	# Pad bit_string to a multiple of 512 chars
	bit_string += b"1"
	while len(__magic_name__					) % 512 != 448:
		bit_string += b"0"
	bit_string += to_little_endian(start_len[32:]					) + to_little_endian(start_len[:32]					)

	return bit_string

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			Generator[list[int], None, None]:






	"""simple docstring"""



	if len(__magic_name__					) % 512 != 0:
		raise ValueError("""Input must have length that's a multiple of 512"""					)

	for pos in range(0	,			len(__magic_name__					)	,			512					):
		UpperCamelCase     :Tuple          =	bit_string[pos : pos + 512]
		UpperCamelCase     :Optional[int]          =	[]
		for i in range(0	,			512	,			32					):
			block_words.append(int(to_little_endian(block[i : i + 32]					)	,			2					)					)
		yield block_words

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int					)		->			int:






	"""simple docstring"""



	if i < 0:
		raise ValueError("""Input must be non-negative"""					)

	UpperCamelCase     :List[str]          =	format(__magic_name__	,			"""032b"""					)
	UpperCamelCase     :Any          =	""""""
	for c in i_str:
		new_str += "1" if c == "0" else "0"
	return int(__magic_name__	,			2					)

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int	,			__magic_name__					:  int					)		->			int:






	"""simple docstring"""



	return (a + b) % 2**32

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int	,			__magic_name__					:  int					)		->			int:






	"""simple docstring"""



	if i < 0:
		raise ValueError("""Input must be non-negative"""					)
	if shift < 0:
		raise ValueError("""Shift must be non-negative"""					)
	return ((i << shift) ^ (i >> (32 - shift))) % 2**32

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			bytes:






	"""simple docstring"""



	UpperCamelCase     :Tuple          =	preprocess(__magic_name__					)

	UpperCamelCase     :List[str]          =	[int(2**32 * abs(sin(i + 1					)					)					) for i in range(64					)]

	# Starting states
	UpperCamelCase     :Union[str, Any]          =	0X67_45_23_01
	UpperCamelCase     :Union[str, Any]          =	0XEF_CD_AB_89
	UpperCamelCase     :List[str]          =	0X98_BA_DC_FE
	UpperCamelCase     :int          =	0X10_32_54_76

	UpperCamelCase     :int          =	[
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	]

	# Process bit string in chunks, each with 16 32-char words
	for block_words in get_block_words(__magic_name__					):
		UpperCamelCase     :Optional[Any]          =	aa
		UpperCamelCase     :Any          =	ba
		UpperCamelCase     :Tuple          =	ca
		UpperCamelCase     :List[str]          =	da

		# Hash current chunk
		for i in range(64					):
			if i <= 15:
				# f = (b & c) | (not_32(b) & d)     # Alternate definition for f
				UpperCamelCase     :int          =	d ^ (b & (c ^ d))
				UpperCamelCase     :Optional[int]          =	i
			elif i <= 31:
				# f = (d & b) | (not_32(d) & c)     # Alternate definition for f
				UpperCamelCase     :str          =	c ^ (d & (b ^ c))
				UpperCamelCase     :Union[str, Any]          =	(5 * i + 1) % 16
			elif i <= 47:
				UpperCamelCase     :str          =	b ^ c ^ d
				UpperCamelCase     :Optional[int]          =	(3 * i + 5) % 16
			else:
				UpperCamelCase     :List[str]          =	c ^ (b | not_aa(__magic_name__					))
				UpperCamelCase     :int          =	(7 * i) % 16
			UpperCamelCase     :Dict          =	(f + a + added_consts[i] + block_words[g]) % 2**32
			UpperCamelCase     :Tuple          =	d
			UpperCamelCase     :str          =	c
			UpperCamelCase     :Tuple          =	b
			UpperCamelCase     :Optional[Any]          =	sum_aa(__magic_name__	,			left_rotate_aa(__magic_name__	,			shift_amounts[i]					)					)

		# Add hashed chunk to running total
		UpperCamelCase     :List[str]          =	sum_aa(__magic_name__	,			__magic_name__					)
		UpperCamelCase     :str          =	sum_aa(__magic_name__	,			__magic_name__					)
		UpperCamelCase     :int          =	sum_aa(__magic_name__	,			__magic_name__					)
		UpperCamelCase     :Optional[Any]          =	sum_aa(__magic_name__	,			__magic_name__					)

	UpperCamelCase     :Optional[Any]          =	reformat_hex(__magic_name__					) + reformat_hex(__magic_name__					) + reformat_hex(__magic_name__					) + reformat_hex(__magic_name__					)
	return digest


if __name__ == "__main__":
	import doctest

	doctest.testmod()
