code: 
  
  
  
  
  
  import posixpath
  from pathlib import Path
  from unittest.mock import patch
  
  import pytest
  from fsspec.implementations.local import AbstractFileSystem, LocalFileSystem, stringify_path
  from fsspec.registry import _registry as _fsspec_registry
  
  class 				A    (__UpperCAmelCase				):
        _SCREAMING_SNAKE_CASE					=							"""mock"""
  
  
  
        def __init__(					self     ,  *lowercase_     ,  lowercase_     ,  **lowercase_					)      ->   Dict:
  
  
  
               '''simple docstring'''
  
               super().__init__()
               _snake_case      :							List[str]     				=						LocalFileSystem(*lowercase_     ,  **lowercase_					)
               _snake_case      :							Any     				=						Path(lowercase_					).resolve().as_posix() + '''/'''
  
  
  
        def    __a (					self     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)      ->   int:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							Dict     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.mkdir(lowercase_     ,  *lowercase_     ,  **lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)      ->   int:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							Union[str, Any]     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.makedirs(lowercase_     ,  *lowercase_     ,  **lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_					)      ->   Union[str, Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							Any     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.rmdir(lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_     ,  lowercase_=True     ,  *lowercase_     ,  **lowercase_					)      ->   List[Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							List[Any]     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               _snake_case      :							Optional[Any]     				=						self._fs.ls(lowercase_     ,  detail=lowercase_     ,  *lowercase_     ,  **lowercase_					)
               if detail:
                      return [{**info, "name": info["name"][len(self.local_root_dir					) :]} for info in out]
               else:
                      return [name[len(self.local_root_dir					) :] for name in out]
  
  
  
        def    __a (					self     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)      ->   Union[str, Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							List[Any]     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               _snake_case      :							Optional[int]     				=						dict(self._fs.info(lowercase_     ,  *lowercase_     ,  **lowercase_					)					)
               _snake_case      :							Union[str, Any]     				=						out['''name'''][len(self.local_root_dir					) :]
               return out
  
  
  
        def    __a (					self     ,  lowercase_     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)      ->   Any:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							Tuple     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               _snake_case      :							int     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.cp_file(lowercase_     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)      ->   List[Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							Tuple     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.rm_file(lowercase_     ,  *lowercase_     ,  **lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)      ->   List[Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							int     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.rm(lowercase_     ,  *lowercase_     ,  **lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_     ,  *lowercase_     ,  **lowercase_					)      ->   Union[str, Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							List[str]     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs._open(lowercase_     ,  *lowercase_     ,  **lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_					)      ->   Optional[Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							Union[str, Any]     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.created(lowercase_					)
  
  
  
        def    __a (					self     ,  lowercase_					)      ->   List[str]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							int     				=						posixpath.join(self.local_root_dir     ,  self._strip_protocol(lowercase_					)					)
               return self._fs.modified(lowercase_					)
  
  
  
        @classmethod
        def    __a (					cls     ,  lowercase_					)      ->   Optional[Any]:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							str     				=						stringify_path(lowercase_					)
               if path.startswith('''mock://'''					):
                      _snake_case      :							int     				=						path[7:]
               return path
  
  
  class 				A    (__UpperCAmelCase				):
        _SCREAMING_SNAKE_CASE					=							"""tmp"""
        _SCREAMING_SNAKE_CASE					=							None
  
  
  
        def __init__(					self     ,  *lowercase_     ,  **lowercase_					)      ->   Dict:
  
  
  
               '''simple docstring'''
  
               assert self.tmp_dir is not None, "TmpDirFileSystem.tmp_dir is not set"
               super().__init__(*lowercase_     ,  **lowercase_     ,  local_root_dir=self.tmp_dir     ,  auto_mkdir=lowercase_					)
  
  
  
        @classmethod
        def    __a (					cls     ,  lowercase_					)      ->   Dict:
  
  
  
               '''simple docstring'''
  
               _snake_case      :							List[str]     				=						stringify_path(lowercase_					)
               if path.startswith('''tmp://'''					):
                      _snake_case      :							int     				=						path[6:]
               return path
  
  
  
  
  @pytest.fixture
  def 						A_			(  )			->						Tuple:
         _snake_case      :							Tuple     				=						MockFileSystem
         _snake_case      :							int     				=						TmpDirFileSystem
         yield
         del _fsspec_registry["mock"]
         del _fsspec_registry["tmp"]
  
  
  
  
  @pytest.fixture
  def 						A_			(  lowercase_      ,		lowercase_  )			->						str:
         _snake_case      :							Tuple     				=						tmp_path_factory.mktemp('''mockfs'''  )
         return MockFileSystem(local_root_dir=lowercase_      ,		auto_mkdir=lowercase_  )
  
  
  
  
  @pytest.fixture
  def 						A_			(  lowercase_      ,		lowercase_  )			->						List[Any]:
         _snake_case      :							Union[str, Any]     				=						tmp_path_factory.mktemp('''tmpfs'''  )
         with patch.object(lowercase_      ,		'''tmp_dir'''      ,		lowercase_  ):
                yield TmpDirFileSystem()
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  from functools import partial
  from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload
  
  import torch
  import torch.nn as nn
  import torch.types
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						torch.Tensor:
         # The first operation in a checkpoint can't be in-place, but it's
         # nice to have in-place addition during inference. Thus...
         if not inplace:
                _snake_case      :							Any     				=						ma + ma
         else:
                ma += ma
  
         return ma
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_  )			->						torch.Tensor:
         _snake_case      :							str     				=						-1 * len(lowercase_  )
         _snake_case      :							Any     				=						list(range(len(tensor.shape[:zero_index]  )  )  )
         return tensor.permute(first_inds + [zero_index + i for i in inds]  )
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_  )			->						torch.Tensor:
         return t.reshape(t.shape[:-no_dims] + (-1,)  )
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_      ,		lowercase_ = 1E-4  )			->						torch.Tensor:
         _snake_case      :							Any     				=						mask.expand(*value.shape  )
         return torch.sum(mask * value      ,		dim=lowercase_  ) / (eps + torch.sum(lowercase_      ,		dim=lowercase_  ))
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_ = 2.31_25      ,		lowercase_ = 21.68_75      ,		lowercase_ = 64  )			->						torch.Tensor:
         _snake_case      :							int     				=						torch.linspace(lowercase_      ,		lowercase_      ,		no_bins - 1      ,		device=pts.device  )
         _snake_case      :							Optional[Any]     				=						torch.sqrt(torch.sum((pts.unsqueeze(-2  ) - pts.unsqueeze(-3  )) ** 2      ,		dim=-1  )  )
         return torch.bucketize(lowercase_      ,		lowercase_  )
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_  )			->						dict:
         _snake_case      :							Dict     				=						dicts[0]
         _snake_case      :							str     				=						{}
         for k, v in first.items():
                _snake_case      :							List[Any]     				=						[d[k] for d in dicts]
                if isinstance(lowercase_      ,		lowercase_  ):
                       _snake_case      :							Any     				=						dict_multimap(lowercase_      ,		lowercase_  )
                else:
                       _snake_case      :							Dict     				=						fn(lowercase_  )
  
         return new_dict
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_  )			->						torch.Tensor:
         _snake_case      :							str     				=						v_bins.view(((1,) * len(x.shape  )) + (len(lowercase_  ),)  )
         _snake_case      :							Union[str, Any]     				=						x[..., None] - reshaped_bins
         _snake_case      :							Optional[Any]     				=						torch.argmin(torch.abs(lowercase_  )      ,		dim=-1  )
         return nn.functional.one_hot(lowercase_      ,		num_classes=len(lowercase_  )  ).float()
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_ = 0      ,		lowercase_ = 0  )			->						torch.Tensor:
         _snake_case      :							List[Union[slice, torch.Tensor]]     				=						[]
         for i, s in enumerate(data.shape[:no_batch_dims]  ):
                _snake_case      :							Optional[Any]     				=						torch.arange(lowercase_  )
                _snake_case      :							Union[str, Any]     				=						r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape  ) - i - 1)))  )
                ranges.append(lowercase_  )
  
         _snake_case      :							List[Union[slice, torch.Tensor]]     				=						[slice(lowercase_  ) for _ in range(len(data.shape  ) - no_batch_dims  )]
         _snake_case      :							Union[str, Any]     				=						inds
         ranges.extend(lowercase_  )
         # Matt note: Editing this to get around the behaviour of using a list as an array index changing
         # in recent Numpy versions
         return data[tuple(lowercase_  )]
  
  
  lowerCAmelCase_   =   TypeVar("T")
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						Dict[Any, Union[dict, list, tuple, Any]]:
         _snake_case      :							Dict[Any, Union[dict, list, tuple, Any]]     				=						{}
         for k, v in dic.items():
                if isinstance(lowercase_      ,		lowercase_  ):
                       _snake_case      :							Dict     				=						dict_map(lowercase_      ,		lowercase_      ,		lowercase_  )
                else:
                       _snake_case      :							List[Any]     				=						tree_map(lowercase_      ,		lowercase_      ,		lowercase_  )
  
         return new_dict
  
  
  
  
  @overload
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						Any:
         ...
  
  
  
  
  @overload
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						dict:
         ...
  
  
  
  
  @overload
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						list:
         ...
  
  
  
  
  @overload
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						tuple:
         ...
  
  
  
  
  def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						Any:
         if isinstance(lowercase_      ,		lowercase_  ):
                return dict_map(lowercase_      ,		lowercase_      ,		lowercase_  )
         elif isinstance(lowercase_      ,		lowercase_  ):
                return [tree_map(lowercase_      ,		lowercase_      ,		lowercase_  ) for x in tree]
         elif isinstance(lowercase_      ,		lowercase_  ):
                return tuple(tree_map(lowercase_      ,		lowercase_      ,		lowercase_  ) for x in tree  )
         elif isinstance(lowercase_      ,		lowercase_  ):
                return fn(lowercase_  )
         else:
                print(type(lowercase_  )  )
                raise ValueError('''Not supported'''  )
  lowerCAmelCase_   =   partial(tree_map, leaf_type=torch.Tensor)
  
  
  
  
  
  
  