






'''simple docstring'''




import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_url
from PIL import Image

from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase__          =      logging.get_logger(__name__)







def       __UpperCAmelCase						(  lowerCamelCase_)							->     Any:
		UpperCamelCase__				:					Dict					=					DPTConfig()

		if "large" in checkpoint_url:
				UpperCamelCase__				:					List[str]					=					1_024
				UpperCamelCase__				:					List[str]					=					4_096
				UpperCamelCase__				:					Optional[int]					=					24
				UpperCamelCase__				:					List[str]					=					16
				UpperCamelCase__				:					List[str]					=					[5, 11, 17, 23]
				UpperCamelCase__				:					str					=					[256, 512, 1_024, 1_024]
				UpperCamelCase__				:					Union[str, Any]					=					(1, 384, 384)

		if "ade" in checkpoint_url:
				UpperCamelCase__				:					int					=					True

				UpperCamelCase__				:					Optional[Any]					=					150
				UpperCamelCase__				:					int					=					'huggingface/label-files'
				UpperCamelCase__				:					List[Any]					=					'ade20k-id2label.json'
				UpperCamelCase__				:					List[Any]					=					json.load(open(cached_download(hf_hub_url(lowerCamelCase_     ,      lowerCamelCase_     ,      repo_type='dataset'))     ,      'r'))
				UpperCamelCase__				:					int					=					{int(lowerCamelCase_): v for k, v in idalabel.items()}
				UpperCamelCase__				:					Union[str, Any]					=					idalabel
				UpperCamelCase__				:					List[str]					=					{v: k for k, v in idalabel.items()}
				UpperCamelCase__				:					Any					=					[1, 150, 480, 480]

		return config, expected_shape







def       __UpperCAmelCase						(  lowerCamelCase_)							->     Optional[Any]:
		UpperCamelCase__				:					Tuple					=					['pretrained.model.head.weight', 'pretrained.model.head.bias']
		for k in ignore_keys:
				state_dict.pop(lowerCamelCase_     ,      lowerCamelCase_)







def       __UpperCAmelCase						(  lowerCamelCase_)							->     Optional[Any]:
		if (
		    "pretrained.model" in name
		    and "cls_token" not in name
		    and "pos_embed" not in name
		    and "patch_embed" not in name
		):
				UpperCamelCase__				:					Union[str, Any]					=					name.replace('pretrained.model'     ,      'dpt.encoder')
		if "pretrained.model" in name:
				UpperCamelCase__				:					Dict					=					name.replace('pretrained.model'     ,      'dpt.embeddings')
		if "patch_embed" in name:
				UpperCamelCase__				:					Tuple					=					name.replace('patch_embed'     ,      'patch_embeddings')
		if "pos_embed" in name:
				UpperCamelCase__				:					Optional[Any]					=					name.replace('pos_embed'     ,      'position_embeddings')
		if "attn.proj" in name:
				UpperCamelCase__				:					List[Any]					=					name.replace('attn.proj'     ,      'attention.output.dense')
		if "proj" in name and "project" not in name:
				UpperCamelCase__				:					Optional[Any]					=					name.replace('proj'     ,      'projection')
		if "blocks" in name:
				UpperCamelCase__				:					int					=					name.replace('blocks'     ,      'layer')
		if "mlp.fc1" in name:
				UpperCamelCase__				:					int					=					name.replace('mlp.fc1'     ,      'intermediate.dense')
		if "mlp.fc2" in name:
				UpperCamelCase__				:					Tuple					=					name.replace('mlp.fc2'     ,      'output.dense')
		if "norm1" in name:
				UpperCamelCase__				:					List[Any]					=					name.replace('norm1'     ,      'layernorm_before')
		if "norm2" in name:
				UpperCamelCase__				:					int					=					name.replace('norm2'     ,      'layernorm_after')
		if "scratch.output_conv" in name:
				UpperCamelCase__				:					Union[str, Any]					=					name.replace('scratch.output_conv'     ,      'head')
		if "scratch" in name:
				UpperCamelCase__				:					int					=					name.replace('scratch'     ,      'neck')
		if "layer1_rn" in name:
				UpperCamelCase__				:					Optional[Any]					=					name.replace('layer1_rn'     ,      'convs.0')
		if "layer2_rn" in name:
				UpperCamelCase__				:					List[Any]					=					name.replace('layer2_rn'     ,      'convs.1')
		if "layer3_rn" in name:
				UpperCamelCase__				:					List[Any]					=					name.replace('layer3_rn'     ,      'convs.2')
		if "layer4_rn" in name:
				UpperCamelCase__				:					List[str]					=					name.replace('layer4_rn'     ,      'convs.3')
		if "refinenet" in name:
				UpperCamelCase__				:					int					=					int(name[len('neck.refinenet') : len('neck.refinenet') + 1])
				# tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3
				UpperCamelCase__				:					Any					=					name.replace(f'refinenet{layer_idx}'     ,      f'fusion_stage.layers.{abs(layer_idx-4)}')
		if "out_conv" in name:
				UpperCamelCase__				:					Union[str, Any]					=					name.replace('out_conv'     ,      'projection')
		if "resConfUnit1" in name:
				UpperCamelCase__				:					int					=					name.replace('resConfUnit1'     ,      'residual_layer1')
		if "resConfUnit2" in name:
				UpperCamelCase__				:					Optional[Any]					=					name.replace('resConfUnit2'     ,      'residual_layer2')
		if "conv1" in name:
				UpperCamelCase__				:					Optional[Any]					=					name.replace('conv1'     ,      'convolution1')
		if "conv2" in name:
				UpperCamelCase__				:					int					=					name.replace('conv2'     ,      'convolution2')
		# readout blocks
		if "pretrained.act_postprocess1.0.project.0" in name:
				UpperCamelCase__				:					Any					=					name.replace('pretrained.act_postprocess1.0.project.0'     ,      'neck.reassemble_stage.readout_projects.0.0')
		if "pretrained.act_postprocess2.0.project.0" in name:
				UpperCamelCase__				:					Tuple					=					name.replace('pretrained.act_postprocess2.0.project.0'     ,      'neck.reassemble_stage.readout_projects.1.0')
		if "pretrained.act_postprocess3.0.project.0" in name:
				UpperCamelCase__				:					int					=					name.replace('pretrained.act_postprocess3.0.project.0'     ,      'neck.reassemble_stage.readout_projects.2.0')
		if "pretrained.act_postprocess4.0.project.0" in name:
				UpperCamelCase__				:					int					=					name.replace('pretrained.act_postprocess4.0.project.0'     ,      'neck.reassemble_stage.readout_projects.3.0')
		# resize blocks
		if "pretrained.act_postprocess1.3" in name:
				UpperCamelCase__				:					Tuple					=					name.replace('pretrained.act_postprocess1.3'     ,      'neck.reassemble_stage.layers.0.projection')
		if "pretrained.act_postprocess1.4" in name:
				UpperCamelCase__				:					Optional[Any]					=					name.replace('pretrained.act_postprocess1.4'     ,      'neck.reassemble_stage.layers.0.resize')
		if "pretrained.act_postprocess2.3" in name:
				UpperCamelCase__				:					Union[str, Any]					=					name.replace('pretrained.act_postprocess2.3'     ,      'neck.reassemble_stage.layers.1.projection')
		if "pretrained.act_postprocess2.4" in name:
				UpperCamelCase__				:					Dict					=					name.replace('pretrained.act_postprocess2.4'     ,      'neck.reassemble_stage.layers.1.resize')
		if "pretrained.act_postprocess3.3" in name:
				UpperCamelCase__				:					Any					=					name.replace('pretrained.act_postprocess3.3'     ,      'neck.reassemble_stage.layers.2.projection')
		if "pretrained.act_postprocess4.3" in name:
				UpperCamelCase__				:					List[Any]					=					name.replace('pretrained.act_postprocess4.3'     ,      'neck.reassemble_stage.layers.3.projection')
		if "pretrained.act_postprocess4.4" in name:
				UpperCamelCase__				:					Optional[Any]					=					name.replace('pretrained.act_postprocess4.4'     ,      'neck.reassemble_stage.layers.3.resize')
		if "pretrained" in name:
				UpperCamelCase__				:					List[str]					=					name.replace('pretrained'     ,      'dpt')
		if "bn" in name:
				UpperCamelCase__				:					Tuple					=					name.replace('bn'     ,      'batch_norm')
		if "head" in name:
				UpperCamelCase__				:					Union[str, Any]					=					name.replace('head'     ,      'head.head')
		if "encoder.norm" in name:
				UpperCamelCase__				:					int					=					name.replace('encoder.norm'     ,      'layernorm')
		if "auxlayer" in name:
				UpperCamelCase__				:					Union[str, Any]					=					name.replace('auxlayer'     ,      'auxiliary_head.head')

		return name







def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_)							->     Any:
		for i in range(config.num_hidden_layers):
				# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
				UpperCamelCase__				:					Optional[int]					=					state_dict.pop(f'dpt.encoder.layer.{i}.attn.qkv.weight')
				UpperCamelCase__				:					Any					=					state_dict.pop(f'dpt.encoder.layer.{i}.attn.qkv.bias')
				# next, add query, keys and values (in that order) to the state dict
				UpperCamelCase__				:					List[str]					=					in_proj_weight[: config.hidden_size, :]
				UpperCamelCase__				:					List[Any]					=					in_proj_bias[: config.hidden_size]
				UpperCamelCase__				:					List[Any]					=					in_proj_weight[
				    config.hidden_size : config.hidden_size * 2, :
				]
				UpperCamelCase__				:					List[Any]					=					in_proj_bias[
				    config.hidden_size : config.hidden_size * 2
				]
				UpperCamelCase__				:					List[str]					=					in_proj_weight[
				    -config.hidden_size :, :
				]
				UpperCamelCase__				:					int					=					in_proj_bias[-config.hidden_size :]







def       __UpperCAmelCase						(  )							->     Optional[Any]:
		UpperCamelCase__				:					Tuple					=					'http://images.cocodataset.org/val2017/000000039769.jpg'
		UpperCamelCase__				:					List[Any]					=					Image.open(requests.get(lowerCamelCase_     ,      stream=lowerCamelCase_).raw)
		return im







@torch.no_grad()
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_)							->     Dict:
		UpperCamelCase__				:					Any					=					get_dpt_config(lowerCamelCase_)
		# load original state_dict from URL
		UpperCamelCase__				:					Tuple					=					torch.hub.load_state_dict_from_url(lowerCamelCase_     ,      map_location='cpu')
		# remove certain keys
		remove_ignore_keys_(lowerCamelCase_)
		# rename keys
		for key in state_dict.copy().keys():
				UpperCamelCase__				:					str					=					state_dict.pop(lowerCamelCase_)
				UpperCamelCase__				:					List[str]					=					val
		# read in qkv matrices
		read_in_q_k_v(lowerCamelCase_     ,      lowerCamelCase_)

		# load HuggingFace model
		UpperCamelCase__				:					str					=					DPTForSemanticSegmentation(lowerCamelCase_) if 'ade' in checkpoint_url else DPTForDepthEstimation(lowerCamelCase_)
		model.load_state_dict(lowerCamelCase_)
		model.eval()

		# Check outputs on an image
		UpperCamelCase__				:					Any					=					480 if 'ade' in checkpoint_url else 384
		UpperCamelCase__				:					List[Any]					=					DPTImageProcessor(size=lowerCamelCase_)

		UpperCamelCase__				:					int					=					prepare_img()
		UpperCamelCase__				:					Optional[Any]					=					image_processor(lowerCamelCase_     ,      return_tensors='pt')

		# forward pass
		UpperCamelCase__				:					Any					=					model(**lowerCamelCase_).logits if 'ade' in checkpoint_url else model(**lowerCamelCase_).predicted_depth

		# Assert logits
		UpperCamelCase__				:					Tuple					=					torch.tensor([[6.3_199, 6.3_629, 6.4_148], [6.3_850, 6.3_615, 6.4_166], [6.3_519, 6.3_176, 6.3_575]])
		if "ade" in checkpoint_url:
				UpperCamelCase__				:					List[str]					=					torch.tensor([[4.0_480, 4.2_420, 4.4_360], [4.3_124, 4.5_693, 4.8_261], [4.5_768, 4.8_965, 5.2_163]])
		assert outputs.shape == torch.Size(lowerCamelCase_)
		assert (
		    torch.allclose(outputs[0, 0, :3, :3]     ,      lowerCamelCase_     ,      atol=1e-4)
		    if "ade" in checkpoint_url
		    else torch.allclose(outputs[0, :3, :3]     ,      lowerCamelCase_)
		)

		Path(lowerCamelCase_).mkdir(exist_ok=lowerCamelCase_)
		print(f'Saving model to {pytorch_dump_folder_path}')
		model.save_pretrained(lowerCamelCase_)
		print(f'Saving image processor to {pytorch_dump_folder_path}')
		image_processor.save_pretrained(lowerCamelCase_)

		if push_to_hub:
				print('Pushing model to hub...')
				model.push_to_hub(
				    repo_path_or_name=Path(lowerCamelCase_     ,      lowerCamelCase_)     ,      organization='nielsr'     ,      commit_message='Add model'     ,      use_temp_dir=lowerCamelCase_     ,      )
				image_processor.push_to_hub(
				    repo_path_or_name=Path(lowerCamelCase_     ,      lowerCamelCase_)     ,      organization='nielsr'     ,      commit_message='Add image processor'     ,      use_temp_dir=lowerCamelCase_     ,      )


if __name__ == "__main__":
		lowerCAmelCase__          =      argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--checkpoint_url',
		    default='https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt',
		    type=str,
		    help='URL of the original DPT checkpoint you\'d like to convert.',
		)
		parser.add_argument(
		    '--pytorch_dump_folder_path',
		    default=None,
		    type=str,
		    required=True,
		    help='Path to the output PyTorch model directory.',
		)
		parser.add_argument(
		    '--push_to_hub',
		    action='store_true',
		)
		parser.add_argument(
		    '--model_name',
		    default='dpt-large',
		    type=str,
		    help='Name of the model, in case you\'re pushing to the hub.',
		)

		lowerCAmelCase__          =      parser.parse_args()
		convert_dpt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)












'''simple docstring'''




import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map







def       __UpperCAmelCase						(  lowerCamelCase_)							->     List[Tuple[int, ...]]:
		UpperCamelCase__				:					int					=					[]
		if isinstance(lowerCamelCase_     ,      lowerCamelCase_):
				for v in tree.values():
						shapes.extend(_fetch_dims(lowerCamelCase_))
		elif isinstance(lowerCamelCase_     ,      (list, tuple)):
				for t in tree:
						shapes.extend(_fetch_dims(lowerCamelCase_))
		elif isinstance(lowerCamelCase_     ,      torch.Tensor):
				shapes.append(tree.shape)
		else:
				raise ValueError('Not supported')

		return shapes







@torch.jit.ignore
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_)							->     Tuple[int, ...]:
		UpperCamelCase__				:					int					=					[]
		for d in reversed(lowerCamelCase_):
				idx.append(flat_idx % d)
				UpperCamelCase__				:					Any					=					flat_idx // d

		return tuple(reversed(lowerCamelCase_))







@torch.jit.ignore
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_ = None     ,      lowerCamelCase_ = None     ,      )							->     List[Tuple[slice, ...]]:

		# start_edges and end_edges both indicate whether, starting from any given
		# dimension, the start/end index is at the top/bottom edge of the
		# corresponding tensor, modeled as a tree
		def reduce_edge_list(lowerCamelCase_) -> None:
				UpperCamelCase__				:					Tuple					=					True
				for i in range(len(lowerCamelCase_)):
						UpperCamelCase__				:					List[Any]					=					-1 * (i + 1)
						l[reversed_idx] &= tally
						UpperCamelCase__				:					Optional[Any]					=					l[reversed_idx]

		if start_edges is None:
				UpperCamelCase__				:					int					=					[s == 0 for s in start]
				reduce_edge_list(lowerCamelCase_)
		if end_edges is None:
				UpperCamelCase__				:					List[str]					=					[e == (d - 1) for e, d in zip(lowerCamelCase_     ,      lowerCamelCase_)]
				reduce_edge_list(lowerCamelCase_)

		# Base cases. Either start/end are empty and we're done, or the final,
		# one-dimensional tensor can be simply sliced
		if len(lowerCamelCase_) == 0:
				return [()]
		elif len(lowerCamelCase_) == 1:
				return [(slice(start[0]     ,      end[0] + 1),)]

		UpperCamelCase__				:					List[Tuple[slice, ...]]					=					[]
		UpperCamelCase__				:					List[slice]					=					[]

		# Dimensions common to start and end can be selected directly
		for s, e in zip(lowerCamelCase_     ,      lowerCamelCase_):
				if s == e:
						path_list.append(slice(lowerCamelCase_     ,      s + 1))
				else:
						break

		UpperCamelCase__				:					Tuple[slice, ...]					=					tuple(lowerCamelCase_)
		UpperCamelCase__				:					Dict					=					len(lowerCamelCase_)

		# start == end, and we're done
		if divergence_idx == len(lowerCamelCase_):
				return [path]

		def upper() -> Tuple[Tuple[slice, ...], ...]:
				assert start_edges is not None
				assert end_edges is not None

				UpperCamelCase__				:					str					=					start[divergence_idx]
				return tuple(
				    path + (slice(lowerCamelCase_     ,      sdi + 1),) + s
				    for s in _get_minimal_slice_set(
				        start[divergence_idx + 1 :]     ,      [d - 1 for d in dims[divergence_idx + 1 :]]     ,      dims[divergence_idx + 1 :]     ,      start_edges=start_edges[divergence_idx + 1 :]     ,      end_edges=[True for _ in end_edges[divergence_idx + 1 :]]     ,      ))

		def lower() -> Tuple[Tuple[slice, ...], ...]:
				assert start_edges is not None
				assert end_edges is not None

				UpperCamelCase__				:					Optional[int]					=					end[divergence_idx]
				return tuple(
				    path + (slice(lowerCamelCase_     ,      edi + 1),) + s
				    for s in _get_minimal_slice_set(
				        [0 for _ in start[divergence_idx + 1 :]]     ,      end[divergence_idx + 1 :]     ,      dims[divergence_idx + 1 :]     ,      start_edges=[True for _ in start_edges[divergence_idx + 1 :]]     ,      end_edges=end_edges[divergence_idx + 1 :]     ,      ))

		# If both start and end are at the edges of the subtree rooted at
		# divergence_idx, we can just select the whole subtree at once
		if start_edges[divergence_idx] and end_edges[divergence_idx]:
				slices.append(path + (slice(start[divergence_idx]     ,      end[divergence_idx] + 1),))
		# If just start is at the edge, we can grab almost all of the subtree,
		# treating only the ragged bottom edge as an edge case
		elif start_edges[divergence_idx]:
				slices.append(path + (slice(start[divergence_idx]     ,      end[divergence_idx]),))
				slices.extend(lower())
		# Analogous to the previous case, but the top is ragged this time
		elif end_edges[divergence_idx]:
				slices.extend(upper())
				slices.append(path + (slice(start[divergence_idx] + 1     ,      end[divergence_idx] + 1),))
		# If both sides of the range are ragged, we need to handle both sides
		# separately. If there's contiguous meat in between them, we can index it
		# in one big chunk
		else:
				slices.extend(upper())
				UpperCamelCase__				:					Dict					=					end[divergence_idx] - start[divergence_idx]
				if middle_ground > 1:
						slices.append(path + (slice(start[divergence_idx] + 1     ,      end[divergence_idx]),))
				slices.extend(lower())

		return slices







@torch.jit.ignore
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_)							->     torch.Tensor:
		UpperCamelCase__				:					List[Any]					=					t.shape[:no_batch_dims]
		UpperCamelCase__				:					Optional[int]					=					list(_flat_idx_to_idx(lowerCamelCase_     ,      lowerCamelCase_))
		# _get_minimal_slice_set is inclusive
		UpperCamelCase__				:					Dict					=					list(_flat_idx_to_idx(flat_end - 1     ,      lowerCamelCase_))

		# Get an ordered list of slices to perform
		UpperCamelCase__				:					int					=					_get_minimal_slice_set(
		    lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      )

		UpperCamelCase__				:					List[Any]					=					[t[s] for s in slices]

		return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors])







def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_ = False     ,      lowerCamelCase_ = None     ,      lowerCamelCase_ = False     ,      )							->     Any:
		if not (len(lowerCamelCase_) > 0):
				raise ValueError('Must provide at least one input')

		UpperCamelCase__				:					int					=					[shape[:no_batch_dims] for shape in _fetch_dims(lowerCamelCase_)]
		UpperCamelCase__				:					int					=					tuple([max(lowerCamelCase_) for s in zip(*lowerCamelCase_)])

		def _prep_inputs(lowerCamelCase_) -> torch.Tensor:
				if not low_mem:
						if not sum(t.shape[:no_batch_dims]) == no_batch_dims:
								UpperCamelCase__				:					List[Any]					=					t.expand(orig_batch_dims + t.shape[no_batch_dims:])
						UpperCamelCase__				:					Optional[int]					=					t.reshape(-1     ,      *t.shape[no_batch_dims:])
				else:
						UpperCamelCase__				:					Optional[int]					=					t.expand(orig_batch_dims + t.shape[no_batch_dims:])
				return t

		UpperCamelCase__				:					Dict[str, Any]					=					tensor_tree_map(_prep_inputs     ,      lowerCamelCase_)
		UpperCamelCase__				:					int					=					None
		if _out is not None:
				UpperCamelCase__				:					Optional[int]					=					tensor_tree_map(lambda lowerCamelCase_: t.view([-1] + list(t.shape[no_batch_dims:]))     ,      _out)

		UpperCamelCase__				:					Dict					=					1
		for d in orig_batch_dims:
				flat_batch_dim *= d

		UpperCamelCase__				:					int					=					flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

		def _select_chunk(lowerCamelCase_) -> torch.Tensor:
				return t[i : i + chunk_size] if t.shape[0] != 1 else t

		UpperCamelCase__				:					List[Any]					=					0
		UpperCamelCase__				:					Optional[Any]					=					prepped_outputs
		for _ in range(lowerCamelCase_):
				# Chunk the input
				if not low_mem:
						UpperCamelCase__				:					str					=					_select_chunk
				else:
						UpperCamelCase__				:					List[Any]					=					partial(
						    _chunk_slice     ,      flat_start=lowerCamelCase_     ,      flat_end=min(lowerCamelCase_     ,      i + chunk_size)     ,      no_batch_dims=len(lowerCamelCase_)     ,      )

				UpperCamelCase__				:					Dict[str, Any]					=					tensor_tree_map(lowerCamelCase_     ,      lowerCamelCase_)

				# Run the layer on the chunk
				UpperCamelCase__				:					List[Any]					=					layer(**lowerCamelCase_)

				# Allocate space for the output
				if out is None:
						UpperCamelCase__				:					Optional[int]					=					tensor_tree_map(lambda lowerCamelCase_: t.new_zeros((flat_batch_dim,) + t.shape[1:])     ,      lowerCamelCase_)

				# Put the chunk in its pre-allocated space
				if isinstance(lowerCamelCase_     ,      lowerCamelCase_):

						def assign(lowerCamelCase_     ,      lowerCamelCase_) -> None:
								for k, v in da.items():
										if isinstance(lowerCamelCase_     ,      lowerCamelCase_):
												assign(lowerCamelCase_     ,      da[k])
										else:
												if _add_into_out:
														v[i : i + chunk_size] += da[k]
												else:
														UpperCamelCase__				:					List[str]					=					da[k]

						assign(lowerCamelCase_     ,      lowerCamelCase_)
				elif isinstance(lowerCamelCase_     ,      lowerCamelCase_):
						for xa, xa in zip(lowerCamelCase_     ,      lowerCamelCase_):
								if _add_into_out:
										xa[i : i + chunk_size] += xa
								else:
										UpperCamelCase__				:					int					=					xa
				elif isinstance(lowerCamelCase_     ,      torch.Tensor):
						if _add_into_out:
								out[i : i + chunk_size] += output_chunk
						else:
								UpperCamelCase__				:					Dict					=					output_chunk
				else:
						raise ValueError('Not supported')

				i += chunk_size

		UpperCamelCase__				:					int					=					tensor_tree_map(lambda lowerCamelCase_: t.view(orig_batch_dims + t.shape[1:])     ,      lowerCamelCase_)

		return out







class       __lowercase					:

	def __init__(      self   :						List[str]  ,					UpperCAmelCase_   :						int = 512  ,					):
			UpperCamelCase__				:					str					=					max_chunk_size
			UpperCamelCase__				:					Optional[int]					=					None
			UpperCamelCase__				:					Optional[tuple]					=					None

	def      __UpperCamelCase							(      self   :						str  ,					UpperCAmelCase_   :						Callable  ,					UpperCAmelCase_   :						tuple  ,					UpperCAmelCase_   :						int):
			logging.info('Tuning chunk size...')

			if min_chunk_size >= self.max_chunk_size:
					return min_chunk_size

			UpperCamelCase__				:					List[int]					=					[2**l for l in range(int(math.log(self.max_chunk_size  ,					2)) + 1)]
			UpperCamelCase__				:					List[Any]					=					[c for c in candidates if c > min_chunk_size]
			UpperCamelCase__				:					List[Any]					=					[min_chunk_size] + candidates
			candidates[-1] += 4

			def test_chunk_size(UpperCAmelCase_   :						int) -> bool:
					try:
							with torch.no_grad():
									fn(*UpperCAmelCase_  ,					chunk_size=UpperCAmelCase_)
							return True
					except RuntimeError:
							return False

			UpperCamelCase__				:					Tuple					=					0
			UpperCamelCase__				:					Dict					=					len(UpperCAmelCase_) - 1
			while i > min_viable_chunk_size_index:
					UpperCamelCase__				:					Optional[int]					=					test_chunk_size(candidates[i])
					if not viable:
							UpperCamelCase__				:					Tuple					=					(min_viable_chunk_size_index + i) // 2
					else:
							UpperCamelCase__				:					Optional[int]					=					i
							UpperCamelCase__				:					Dict					=					(i + len(UpperCAmelCase_) - 1) // 2

			return candidates[min_viable_chunk_size_index]

	def      __UpperCamelCase							(      self   :						Any  ,					UpperCAmelCase_   :						Iterable  ,					UpperCAmelCase_   :						Iterable):
			UpperCamelCase__				:					List[str]					=					True
			for aa, aa in zip(UpperCAmelCase_  ,					UpperCAmelCase_):
					assert type(UpperCAmelCase_) == type(UpperCAmelCase_)
					if isinstance(UpperCAmelCase_  ,					(list, tuple)):
							consistent &= self._compare_arg_caches(UpperCAmelCase_  ,					UpperCAmelCase_)
					elif isinstance(UpperCAmelCase_  ,					UpperCAmelCase_):
							UpperCamelCase__				:					Union[str, Any]					=					[v for _, v in sorted(aa.items()  ,					key=lambda UpperCAmelCase_: x[0])]
							UpperCamelCase__				:					str					=					[v for _, v in sorted(aa.items()  ,					key=lambda UpperCAmelCase_: x[0])]
							consistent &= self._compare_arg_caches(UpperCAmelCase_  ,					UpperCAmelCase_)
					else:
							consistent &= aa == aa

			return consistent





	def      __UpperCamelCase							(      self   :						List[Any]  ,					UpperCAmelCase_   :						Callable  ,					UpperCAmelCase_   :						tuple  ,					UpperCAmelCase_   :						int  ,					):
			UpperCamelCase__				:					List[Any]					=					True
			UpperCamelCase__				:					tuple					=					tree_map(lambda UpperCAmelCase_: a.shape if isinstance(UpperCAmelCase_  ,					torch.Tensor) else a  ,					UpperCAmelCase_  ,					UpperCAmelCase_)
			if self.cached_arg_data is not None:
					# If args have changed shape/value, we need to re-tune
					assert len(self.cached_arg_data) == len(UpperCAmelCase_)
					UpperCamelCase__				:					Union[str, Any]					=					self._compare_arg_caches(self.cached_arg_data  ,					UpperCAmelCase_)
			else:
					# Otherwise, we can reuse the precomputed value
					UpperCamelCase__				:					Optional[int]					=					False

			if not consistent:
					UpperCamelCase__				:					Tuple					=					self._determine_favorable_chunk_size(
					    UpperCAmelCase_  ,					UpperCAmelCase_  ,					UpperCAmelCase_  ,					)
					UpperCamelCase__				:					Optional[Any]					=					arg_data

			assert self.cached_chunk_size is not None

			return self.cached_chunk_size




