



'''simple docstring'''

import argparse
import os

import torch

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)


UpperCAmelCase     		=   {
    'sample_size': 32,
    'in_channels': 3,
    'out_channels': 3,
    'layers_per_block': 2,
    'num_class_embeds': 1000,
    'block_out_channels': [32, 64],
    'attention_head_dim': 8,
    'down_block_types': [
        'ResnetDownsampleBlock2D',
        'AttnDownBlock2D',
    ],
    'up_block_types': [
        'AttnUpBlock2D',
        'ResnetUpsampleBlock2D',
    ],
    'resnet_time_scale_shift': 'scale_shift',
    'upsample_type': 'resnet',
    'downsample_type': 'resnet',
}

UpperCAmelCase     		=   {
    'sample_size': 64,
    'in_channels': 3,
    'out_channels': 3,
    'layers_per_block': 3,
    'num_class_embeds': 1000,
    'block_out_channels': [192, 192 * 2, 192 * 3, 192 * 4],
    'attention_head_dim': 64,
    'down_block_types': [
        'ResnetDownsampleBlock2D',
        'AttnDownBlock2D',
        'AttnDownBlock2D',
        'AttnDownBlock2D',
    ],
    'up_block_types': [
        'AttnUpBlock2D',
        'AttnUpBlock2D',
        'AttnUpBlock2D',
        'ResnetUpsampleBlock2D',
    ],
    'resnet_time_scale_shift': 'scale_shift',
    'upsample_type': 'resnet',
    'downsample_type': 'resnet',
}

UpperCAmelCase     		=   {
    'sample_size': 256,
    'in_channels': 3,
    'out_channels': 3,
    'layers_per_block': 2,
    'num_class_embeds': None,
    'block_out_channels': [256, 256, 256 * 2, 256 * 2, 256 * 4, 256 * 4],
    'attention_head_dim': 64,
    'down_block_types': [
        'ResnetDownsampleBlock2D',
        'ResnetDownsampleBlock2D',
        'ResnetDownsampleBlock2D',
        'AttnDownBlock2D',
        'AttnDownBlock2D',
        'AttnDownBlock2D',
    ],
    'up_block_types': [
        'AttnUpBlock2D',
        'AttnUpBlock2D',
        'AttnUpBlock2D',
        'ResnetUpsampleBlock2D',
        'ResnetUpsampleBlock2D',
        'ResnetUpsampleBlock2D',
    ],
    'resnet_time_scale_shift': 'default',
    'upsample_type': 'resnet',
    'downsample_type': 'resnet',
}

UpperCAmelCase     		=   {
    'num_train_timesteps': 40,
    'sigma_min': 0.002,
    'sigma_max': 80.0,
}

UpperCAmelCase     		=   {
    'num_train_timesteps': 201,
    'sigma_min': 0.002,
    'sigma_max': 80.0,
}

UpperCAmelCase     		=   {
    'num_train_timesteps': 151,
    'sigma_min': 0.002,
    'sigma_max': 80.0,
}






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						Optional[Any]							)      ->						Union[str, Any]:







	"""simple docstring"""







	if isinstance(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							):
		return v
	if v.lower() in ("yes", "true", "t", "y", "1"):
		return True
	elif v.lower() in ("no", "false", "f", "n", "0"):
		return False
	else:
		raise argparse.ArgumentTypeError("""boolean value expected"""							)






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						Dict							,		_SCREAMING_SNAKE_CASE					:						Optional[int]							,		_SCREAMING_SNAKE_CASE					:						Optional[int]							,		_SCREAMING_SNAKE_CASE					:						Tuple							,		_SCREAMING_SNAKE_CASE					:						Any=False							)      ->						Optional[int]:







	"""simple docstring"""
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.in_layers.0.weight']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.in_layers.0.bias']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.in_layers.2.weight']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.in_layers.2.bias']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.emb_layers.1.weight']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.emb_layers.1.bias']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.out_layers.0.weight']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.out_layers.0.bias']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.out_layers.3.weight']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.out_layers.3.bias']

	if has_skip:
		lowerCAmelCase   =     checkpoint[f'{old_prefix}.skip_connection.weight']
		lowerCAmelCase   =     checkpoint[f'{old_prefix}.skip_connection.bias']

	return new_checkpoint






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						Dict							,		_SCREAMING_SNAKE_CASE					:						List[Any]							,		_SCREAMING_SNAKE_CASE					:						Optional[int]							,		_SCREAMING_SNAKE_CASE					:						Union[str, Any]							,		_SCREAMING_SNAKE_CASE					:						Optional[int]=None							)      ->						List[str]:







	"""simple docstring"""
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.qkv.weight'].chunk(3							,		dim=0							)
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.qkv.bias'].chunk(3							,		dim=0							)

	lowerCAmelCase   =     checkpoint[f'{old_prefix}.norm.weight']
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.norm.bias']

	lowerCAmelCase   =     weight_q.squeeze(-1							).squeeze(-1							)
	lowerCAmelCase   =     bias_q.squeeze(-1							).squeeze(-1							)
	lowerCAmelCase   =     weight_k.squeeze(-1							).squeeze(-1							)
	lowerCAmelCase   =     bias_k.squeeze(-1							).squeeze(-1							)
	lowerCAmelCase   =     weight_v.squeeze(-1							).squeeze(-1							)
	lowerCAmelCase   =     bias_v.squeeze(-1							).squeeze(-1							)

	lowerCAmelCase   =     (
	    checkpoint[f'{old_prefix}.proj_out.weight'].squeeze(-1							).squeeze(-1							)
	)
	lowerCAmelCase   =     checkpoint[f'{old_prefix}.proj_out.bias'].squeeze(-1							).squeeze(-1							)

	return new_checkpoint






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						Tuple							,		_SCREAMING_SNAKE_CASE					:						List[str]							)      ->						List[Any]:







	"""simple docstring"""
	lowerCAmelCase   =     torch.load(_SCREAMING_SNAKE_CASE							,		map_location="""cpu"""							)
	lowerCAmelCase   =     {}

	lowerCAmelCase   =     checkpoint['''time_embed.0.weight''']
	lowerCAmelCase   =     checkpoint['''time_embed.0.bias''']
	lowerCAmelCase   =     checkpoint['''time_embed.2.weight''']
	lowerCAmelCase   =     checkpoint['''time_embed.2.bias''']

	if unet_config["num_class_embeds"] is not None:
		lowerCAmelCase   =     checkpoint['''label_emb.weight''']

	lowerCAmelCase   =     checkpoint['''input_blocks.0.0.weight''']
	lowerCAmelCase   =     checkpoint['''input_blocks.0.0.bias''']

	lowerCAmelCase   =     unet_config['''down_block_types''']
	lowerCAmelCase   =     unet_config['''layers_per_block''']
	lowerCAmelCase   =     unet_config['''attention_head_dim''']
	lowerCAmelCase   =     unet_config['''block_out_channels''']
	lowerCAmelCase   =     1
	lowerCAmelCase   =     channels_list[0]

	for i, layer_type in enumerate(_SCREAMING_SNAKE_CASE							):
		lowerCAmelCase   =     channels_list[i]
		lowerCAmelCase   =     current_channels != prev_channels
		if layer_type == "ResnetDownsampleBlock2D":
			for j in range(_SCREAMING_SNAKE_CASE							):
				lowerCAmelCase   =     f'down_blocks.{i}.resnets.{j}'
				lowerCAmelCase   =     f'input_blocks.{current_layer}.0'
				lowerCAmelCase   =     True if j == 0 and downsample_block_has_skip else False
				lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		has_skip=_SCREAMING_SNAKE_CASE							)
				current_layer += 1

		elif layer_type == "AttnDownBlock2D":
			for j in range(_SCREAMING_SNAKE_CASE							):
				lowerCAmelCase   =     f'down_blocks.{i}.resnets.{j}'
				lowerCAmelCase   =     f'input_blocks.{current_layer}.0'
				lowerCAmelCase   =     True if j == 0 and downsample_block_has_skip else False
				lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		has_skip=_SCREAMING_SNAKE_CASE							)
				lowerCAmelCase   =     f'down_blocks.{i}.attentions.{j}'
				lowerCAmelCase   =     f'input_blocks.{current_layer}.1'
				lowerCAmelCase   =     convert_attention(
				    _SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)
				current_layer += 1

		if i != len(_SCREAMING_SNAKE_CASE							) - 1:
			lowerCAmelCase   =     f'down_blocks.{i}.downsamplers.0'
			lowerCAmelCase   =     f'input_blocks.{current_layer}.0'
			lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)
			current_layer += 1

		lowerCAmelCase   =     current_channels

	# hardcoded the mid-block for now
	lowerCAmelCase   =     '''mid_block.resnets.0'''
	lowerCAmelCase   =     '''middle_block.0'''
	lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)
	lowerCAmelCase   =     '''mid_block.attentions.0'''
	lowerCAmelCase   =     '''middle_block.1'''
	lowerCAmelCase   =     convert_attention(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)
	lowerCAmelCase   =     '''mid_block.resnets.1'''
	lowerCAmelCase   =     '''middle_block.2'''
	lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)

	lowerCAmelCase   =     0
	lowerCAmelCase   =     unet_config['''up_block_types''']

	for i, layer_type in enumerate(_SCREAMING_SNAKE_CASE							):
		if layer_type == "ResnetUpsampleBlock2D":
			for j in range(layers_per_block + 1							):
				lowerCAmelCase   =     f'up_blocks.{i}.resnets.{j}'
				lowerCAmelCase   =     f'output_blocks.{current_layer}.0'
				lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		has_skip=_SCREAMING_SNAKE_CASE							)
				current_layer += 1

			if i != len(_SCREAMING_SNAKE_CASE							) - 1:
				lowerCAmelCase   =     f'up_blocks.{i}.upsamplers.0'
				lowerCAmelCase   =     f'output_blocks.{current_layer-1}.1'
				lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)
		elif layer_type == "AttnUpBlock2D":
			for j in range(layers_per_block + 1							):
				lowerCAmelCase   =     f'up_blocks.{i}.resnets.{j}'
				lowerCAmelCase   =     f'output_blocks.{current_layer}.0'
				lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		has_skip=_SCREAMING_SNAKE_CASE							)
				lowerCAmelCase   =     f'up_blocks.{i}.attentions.{j}'
				lowerCAmelCase   =     f'output_blocks.{current_layer}.1'
				lowerCAmelCase   =     convert_attention(
				    _SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)
				current_layer += 1

			if i != len(_SCREAMING_SNAKE_CASE							) - 1:
				lowerCAmelCase   =     f'up_blocks.{i}.upsamplers.0'
				lowerCAmelCase   =     f'output_blocks.{current_layer-1}.2'
				lowerCAmelCase   =     convert_resnet(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							)

	lowerCAmelCase   =     checkpoint['''out.0.weight''']
	lowerCAmelCase   =     checkpoint['''out.0.bias''']
	lowerCAmelCase   =     checkpoint['''out.2.weight''']
	lowerCAmelCase   =     checkpoint['''out.2.bias''']

	return new_checkpoint


if __name__ == "__main__":
							UpperCAmelCase     		=   argparse.ArgumentParser()

							parser.add_argument('--unet_path', default=None, type=str, required=True, help='Path to the unet.pt to convert.')
							parser.add_argument(
							    '--dump_path', default=None, type=str, required=True, help='Path to output the converted UNet model.'
							)
							parser.add_argument('--class_cond', default=True, type=str, help='Whether the model is class-conditional.')

							UpperCAmelCase     		=   parser.parse_args()
							UpperCAmelCase     		=   strabool(args.class_cond)

							UpperCAmelCase     		=   os.path.basename(args.unet_path)
							print(F'''Checkpoint: {ckpt_name}''')

							# Get U-Net config
							if "imagenet64" in ckpt_name:
														UpperCAmelCase     		=   IMAGENET_64_UNET_CONFIG
							elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
														UpperCAmelCase     		=   LSUN_256_UNET_CONFIG
							elif "test" in ckpt_name:
														UpperCAmelCase     		=   TEST_UNET_CONFIG
							else:
														raise ValueError(F'''Checkpoint type {ckpt_name} is not currently supported.''')

							if not args.class_cond:
														UpperCAmelCase     		=   None

							UpperCAmelCase     		=   con_pt_to_diffuser(args.unet_path, unet_config)

							UpperCAmelCase     		=   UNetaDModel(**unet_config)
							image_unet.load_state_dict(converted_unet_ckpt)

							# Get scheduler config
							if "cd" in ckpt_name or "test" in ckpt_name:
														UpperCAmelCase     		=   CD_SCHEDULER_CONFIG
							elif "ct" in ckpt_name and "imagenet64" in ckpt_name:
														UpperCAmelCase     		=   CT_IMAGENET_64_SCHEDULER_CONFIG
							elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
														UpperCAmelCase     		=   CT_LSUN_256_SCHEDULER_CONFIG
							else:
														raise ValueError(F'''Checkpoint type {ckpt_name} is not currently supported.''')

							UpperCAmelCase     		=   CMStochasticIterativeScheduler(**scheduler_config)

							UpperCAmelCase     		=   ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler)
							consistency_model.save_pretrained(args.dump_path)




'''simple docstring'''

import tensorflow as tf

from ...tf_utils import shape_list

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


						'''simple docstring'''



						def __init__( self     ,	A_     ,	A_     ,	A_     ,	A_     ,	A_=1     ,	A_=False     ,	**A_				)     ->     Optional[int]:
							super().__init__(**A_				)

							lowerCAmelCase   =     vocab_size
							lowerCAmelCase   =     d_embed
							lowerCAmelCase   =     d_proj

							lowerCAmelCase   =     cutoffs + [vocab_size]
							lowerCAmelCase   =     [0] + self.cutoffs
							lowerCAmelCase   =     div_val

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

							lowerCAmelCase   =     []
							lowerCAmelCase   =     []


						def     __snake_case   ( self     ,	A_				)     ->     int:
							if self.n_clusters > 0:
								lowerCAmelCase   =     self.add_weight(
								    shape=(self.n_clusters, self.d_embed)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name="""cluster_weight"""				)
								lowerCAmelCase   =     self.add_weight(
								    shape=(self.n_clusters,)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name="""cluster_bias"""				)

							if self.div_val == 1:
								for i in range(len(self.cutoffs				)				):
									if self.d_proj != self.d_embed:
										lowerCAmelCase   =     self.add_weight(
										    shape=(self.d_embed, self.d_proj)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_projs_._{i}'     ,	)
										self.out_projs.append(A_				)
									else:
										self.out_projs.append(A_				)
									lowerCAmelCase   =     self.add_weight(
									    shape=(self.vocab_size, self.d_embed)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._weight'     ,	)
									lowerCAmelCase   =     self.add_weight(
									    shape=(self.vocab_size,)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._bias'     ,	)
									self.out_layers.append((weight, bias)				)
							else:
								for i in range(len(self.cutoffs				)				):
									lowerCAmelCase,       lowerCAmelCase   =     self.cutoff_ends[i], self.cutoff_ends[i + 1]
									lowerCAmelCase   =     self.d_embed // (self.div_val**i)

									lowerCAmelCase   =     self.add_weight(
									    shape=(d_emb_i, self.d_proj)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_projs_._{i}'				)
									self.out_projs.append(A_				)
									lowerCAmelCase   =     self.add_weight(
									    shape=(r_idx - l_idx, d_emb_i)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._weight'     ,	)
									lowerCAmelCase   =     self.add_weight(
									    shape=(r_idx - l_idx,)     ,	initializer="""zeros"""     ,	trainable=A_     ,	name=f'out_layers_._{i}_._bias'     ,	)
									self.out_layers.append((weight, bias)				)
							super().build(A_				)


						@staticmethod
						def     __snake_case   ( A_     ,	A_     ,	A_     ,	A_=None				)     ->     List[Any]:
							lowerCAmelCase   =     x
							if proj is not None:
								lowerCAmelCase   =     tf.einsum("""ibd,ed->ibe"""     ,	A_     ,	A_				)
							return tf.einsum("""ibd,nd->ibn"""     ,	A_     ,	A_				) + b


						@staticmethod
						def     __snake_case   ( A_     ,	A_				)     ->     Dict:
							lowerCAmelCase   =     shape_list(A_				)
							lowerCAmelCase   =     tf.range(lp_size[0]     ,	dtype=target.dtype				)
							lowerCAmelCase   =     tf.stack([r, target]     ,	1				)
							return tf.gather_nd(A_     ,	A_				)




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

									if self.div_val == 1:
										lowerCAmelCase   =     self.out_layers[0][0][l_idx:r_idx]
										lowerCAmelCase   =     self.out_layers[0][1][l_idx:r_idx]
									else:
										lowerCAmelCase   =     self.out_layers[i][0]
										lowerCAmelCase   =     self.out_layers[i][1]

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

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

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

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

							return out