code: 
  
  
  
  
  
  from typing import SupportsIndex
  
  import numpy as np
  from scipy.ndimage import convolve
  
  def     a						(      _lowercase   :     np.ndarray ,						_lowercase   :     np.ndarray ,						_lowercase   :     np.ndarray       )     ->np.ndarray:
  		UpperCamelCase__         =	np.stack((horizontal_flow, vertical_flow) ,						2       )
  
  		# Create a grid of all pixel coordinates and subtract the flow to get the
  		# target pixels coordinates
  		UpperCamelCase__         =	np.stack(
  		    np.meshgrid(np.arange(0 ,						image.shape[1]       ) ,						np.arange(0 ,						image.shape[0]       )       ) ,						2       )
  		UpperCamelCase__         =	np.round(grid - flow       ).astype(np.intaa       )
  
  		# Find the locations outside of the original image
  		UpperCamelCase__         =	(grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]       ))
  		UpperCamelCase__         =	0
  
  		UpperCamelCase__         =	image[grid[:, :, 1], grid[:, :, 0]]
  
  		# Set pixels at invalid locations to 0
  		UpperCamelCase__         =	0
  
  		return warped
  
  def     a						(      _lowercase   :     np.ndarray ,						_lowercase   :     np.ndarray ,						_lowercase   :     SupportsIndex ,						_lowercase   :     float | None = None ,						)     ->tuple[np.ndarray, np.ndarray]:
  		if alpha is None:
  				UpperCamelCase__         =	0.1
  
  		# Initialize flow
  		UpperCamelCase__         =	np.zeros_like(_lowercase       )
  		UpperCamelCase__         =	np.zeros_like(_lowercase       )
  
  		# Prepare kernels for the calculation of the derivatives and the average velocity
  		UpperCamelCase__         =	np.array([[-1, 1], [-1, 1]]       ) * 0.25
  		UpperCamelCase__         =	np.array([[-1, -1], [1, 1]]       ) * 0.25
  		UpperCamelCase__         =	np.array([[1, 1], [1, 1]]       ) * 0.25
  		UpperCamelCase__         =	np.array(
  		    [[1 / 1_2, 1 / 6, 1 / 1_2], [1 / 6, 0, 1 / 6], [1 / 1_2, 1 / 6, 1 / 1_2]]       )
  
  		# Iteratively refine the flow
  		for _ in range(_lowercase       ):
  				UpperCamelCase__         =	warp(_lowercase ,						_lowercase ,						_lowercase       )
  				UpperCamelCase__         =	convolve(_lowercase ,						_lowercase       ) + convolve(_lowercase ,						_lowercase       )
  				UpperCamelCase__         =	convolve(_lowercase ,						_lowercase       ) + convolve(_lowercase ,						_lowercase       )
  				UpperCamelCase__         =	convolve(_lowercase ,						_lowercase       ) + convolve(_lowercase ,						-kernel_t       )
  
  				UpperCamelCase__         =	convolve(_lowercase ,						_lowercase       )
  				UpperCamelCase__         =	convolve(_lowercase ,						_lowercase       )
  
  				# This updates the flow as proposed in the paper (Step 12)
  				UpperCamelCase__         =	(
  				    derivative_x * avg_horizontal_velocity
  				    + derivative_y * avg_vertical_velocity
  				    + derivative_t
  				)
  				UpperCamelCase__         =	update / (alpha**2 + derivative_x**2 + derivative_y**2)
  
  				UpperCamelCase__         =	avg_horizontal_velocity - derivative_x * update
  				UpperCamelCase__         =	avg_vertical_velocity - derivative_y * update
  
  		return horizontal_flow, vertical_flow
  
  
  if __name__ == "__main__":
  				import doctest
  
  				doctest.testmod()
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  import warnings
  
  from .generation import FlaxGenerationMixin
  
  
  
  
  class   _UpperCAmelCase	(   _lowercase					):
  					# warning at import time
  					warnings.warn(
  					    "Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will "
  					    "be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead." ,       _lowercase ,       )
  
  
  
  
  
  