






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


a     :   Optional[Any]    =       {
    'configuration_marian': ['MARIAN_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MarianConfig', 'MarianOnnxConfig'],
}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a     :   Any    =       ['MarianTokenizer']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a     :   int    =       [
     'MARIAN_PRETRAINED_MODEL_ARCHIVE_LIST',
     'MarianForCausalLM',
     'MarianModel',
     'MarianMTModel',
     'MarianPreTrainedModel',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a     :   Union[str, Any]    =       ['TFMarianModel', 'TFMarianMTModel', 'TFMarianPreTrainedModel']

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a     :   Union[str, Any]    =       ['FlaxMarianModel', 'FlaxMarianMTModel', 'FlaxMarianPreTrainedModel']

if TYPE_CHECKING:
 from .configuration_marian import MARIAN_PRETRAINED_CONFIG_ARCHIVE_MAP, MarianConfig, MarianOnnxConfig

 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_marian import MarianTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_marian import (
      MARIAN_PRETRAINED_MODEL_ARCHIVE_LIST,
      MarianForCausalLM,
      MarianModel,
      MarianMTModel,
      MarianPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_marian import TFMarianModel, TFMarianMTModel, TFMarianPreTrainedModel

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_marian import FlaxMarianModel, FlaxMarianMTModel, FlaxMarianPreTrainedModel

else:
 import sys

 a     :   Optional[int]    =       _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)










from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve
def       a__						(  A		,				A		,				A					):
   lowercase__						:	Dict         = np.stack((horizontal_flow, vertical_flow)		,				2					)

   # Create a grid of all pixel coordinates and subtract the flow to get the
   # target pixels coordinates
   lowercase__						:	Tuple         = np.stack(
       np.meshgrid(np.arange(0		,				image.shape[1]					)		,				np.arange(0		,				image.shape[0]					)					)		,				2					)
   lowercase__						:	Union[str, Any]         = np.round(grid - flow					).astype(np.intaa					)

   # Find the locations outside of the original image
   lowercase__						:	Optional[Any]         = (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]					))
   lowercase__						:	int         = 0

   lowercase__						:	Tuple         = image[grid[:, :, 1], grid[:, :, 0]]

   # Set pixels at invalid locations to 0
   lowercase__						:	str         = 0

   return warped
def       a__						(  A		,				A		,				A		,				A = None		,				):
   if alpha is None:
      lowercase__						:	Tuple         = 0.1

   # Initialize flow
   lowercase__						:	Optional[int]         = np.zeros_like(A					)
   lowercase__						:	Tuple         = np.zeros_like(A					)

   # Prepare kernels for the calculation of the derivatives and the average velocity
   lowercase__						:	Tuple         = np.array([[-1, 1], [-1, 1]]					) * 0.2_5
   lowercase__						:	int         = np.array([[-1, -1], [1, 1]]					) * 0.2_5
   lowercase__						:	Optional[int]         = np.array([[1, 1], [1, 1]]					) * 0.2_5
   lowercase__						:	Optional[Any]         = np.array(
       [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]					)

   # Iteratively refine the flow
   for _ in range(A					):
      lowercase__						:	Optional[int]         = warp(A		,				A		,				A					)
      lowercase__						:	Union[str, Any]         = convolve(A		,				A					) + convolve(A		,				A					)
      lowercase__						:	Union[str, Any]         = convolve(A		,				A					) + convolve(A		,				A					)
      lowercase__						:	str         = convolve(A		,				A					) + convolve(A		,				-kernel_t					)

      lowercase__						:	str         = convolve(A		,				A					)
      lowercase__						:	Tuple         = convolve(A		,				A					)

      # This updates the flow as proposed in the paper (Step 12)
      lowercase__						:	List[Any]         = (
          derivative_x * avg_horizontal_velocity
          + derivative_y * avg_vertical_velocity
          + derivative_t
      )
      lowercase__						:	Union[str, Any]         = update / (alpha**2 + derivative_x**2 + derivative_y**2)

      lowercase__						:	Optional[int]         = avg_horizontal_velocity - derivative_x * update
      lowercase__						:	int         = avg_vertical_velocity - derivative_y * update

   return horizontal_flow, vertical_flow


if __name__ == "__main__":
 import doctest

 doctest.testmod()


