from __future__ import annotations

__lowerCamelCase			: List[str]							   =		list[list[int]]

# assigning initial values to the grid
__lowerCamelCase			: Matrix							   =		[
    [3, 0, 6, 5, 0, 8, 4, 0, 0],
    [5, 2, 0, 0, 0, 0, 0, 0, 0],
    [0, 8, 7, 0, 0, 0, 0, 3, 1],
    [0, 0, 3, 0, 1, 0, 0, 8, 0],
    [9, 0, 0, 8, 6, 3, 0, 0, 5],
    [0, 5, 0, 0, 9, 0, 6, 0, 0],
    [1, 3, 0, 0, 0, 0, 2, 5, 0],
    [0, 0, 0, 0, 0, 0, 0, 7, 4],
    [0, 0, 5, 2, 0, 6, 3, 0, 0],
]

# a grid with no solution
__lowerCamelCase			: Matrix							   =		[
    [5, 0, 6, 5, 0, 8, 4, 0, 3],
    [5, 2, 0, 0, 0, 0, 0, 0, 2],
    [1, 8, 7, 0, 0, 0, 0, 3, 1],
    [0, 0, 3, 0, 1, 0, 0, 8, 0],
    [9, 0, 0, 8, 6, 3, 0, 0, 5],
    [0, 5, 0, 0, 9, 0, 6, 0, 0],
    [1, 3, 0, 0, 0, 0, 2, 5, 0],
    [0, 0, 0, 0, 0, 0, 0, 7, 4],
    [0, 0, 5, 2, 0, 6, 3, 0, 0],
]



def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       Matrix ,     __UpperCamelCase    :       int ,     __UpperCamelCase    :       int ,     __UpperCamelCase    :       int	)							->					bool:







  """simple docstring"""





  for i in range(9	):
    if grid[row][i] == n or grid[i][column] == n:
      return False

  for i in range(3	):
    for j in range(3	):
      if grid[(row - row % 3) + i][(column - column % 3) + j] == n:
        return False

  return True



def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       Matrix	)							->					tuple[int, int] | None:







  """simple docstring"""





  for i in range(9	):
    for j in range(9	):
      if grid[i][j] == 0:
        return i, j
  return None



def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       Matrix	)							->					Matrix | None:







  """simple docstring"""





  if location := find_empty_location(__UpperCamelCase	):
    SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      location
  else:
    # If the location is ``None``, then the grid is solved.
    return grid

  for digit in range(1 ,     10	):
    if is_safe(__UpperCamelCase ,     __UpperCamelCase ,     __UpperCamelCase ,     __UpperCamelCase	):
      SCREAMING_SNAKE_CASE__										=      digit

      if sudoku(__UpperCamelCase	) is not None:
        return grid

      SCREAMING_SNAKE_CASE__										=      0

  return None



def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       Matrix	)							->					None:







  """simple docstring"""





  for row in grid:
    for cell in row:
      print(__UpperCamelCase ,     end=""" """	)
    print()


if __name__ == "__main__":
   # make a copy of grid so that you can compare with the unmodified grid
   for example_grid in (initial_grid, no_solution):
      print('''\nExample grid:\n''' + '''=''' * 20)
      print_solution(example_grid)
      print('''\nExample grid solution:''')
      __lowerCamelCase			: Optional[Any]							   =		sudoku(example_grid)
      if solution is not None:
         print_solution(solution)
      else:
         print('''Cannot find a solution.''')

import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__lowerCamelCase			: List[str]							   =		object()

# For specifying empty leaf dict `{}`
__lowerCamelCase			: Optional[int]							   =		object()



def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       List[Any] ,     __UpperCamelCase    :       List[str]	)							->					Optional[Any]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__										=      tuple((re.compile(x + """$"""	) for x in qs)	)
  for i in range(len(__UpperCamelCase	) - len(__UpperCamelCase	) + 1	):
    SCREAMING_SNAKE_CASE__										=      [x.match(__UpperCamelCase	) for x, y in zip(__UpperCamelCase ,     ks[i:]	)]
    if matches and all(__UpperCamelCase	):
      return True
  return False



def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       Optional[int]	)							->					List[str]:







  """simple docstring"""





  def replace(__UpperCamelCase    :       Tuple ,     __UpperCamelCase    :       Any	):
    for rule, replacement in rules:
      if _match(__UpperCamelCase ,     __UpperCamelCase	):
        return replacement
    return val

  return replace



def 	__SCREAMING_SNAKE_CASE  (					)							->					List[Any]:







  """simple docstring"""





  return [
      # embeddings
      (("transformer", "wpe", "embedding"), P("""mp""" ,     __UpperCamelCase	)),
      (("transformer", "wte", "embedding"), P("""mp""" ,     __UpperCamelCase	)),
      # atention
      (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__UpperCamelCase ,     """mp"""	)),
      (("attention", "out_proj", "kernel"), P("""mp""" ,     __UpperCamelCase	)),
      (("attention", "out_proj", "bias"), None),
      # mlp
      (("mlp", "c_fc", "kernel"), P(__UpperCamelCase ,     """mp"""	)),
      (("mlp", "c_fc", "bias"), P("""mp"""	)),
      (("mlp", "c_proj", "kernel"), P("""mp""" ,     __UpperCamelCase	)),
      (("mlp", "c_proj", "bias"), None),
      # layer norms
      ((r"ln_\d+", "bias"), None),
      ((r"\d+", r"ln_\d+", "scale"), None),
      (("ln_f", "bias"), None),
      (("ln_f", "scale"), None),
  ]






def 	__SCREAMING_SNAKE_CASE  (					__UpperCamelCase    :       Union[str, Any]	)							->					Dict:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__										=      _get_partition_rules()
  SCREAMING_SNAKE_CASE__										=      _replacement_rules(__UpperCamelCase	)
  SCREAMING_SNAKE_CASE__										=      {k: _unmatched for k in flatten_dict(__UpperCamelCase	)}
  SCREAMING_SNAKE_CASE__										=      {k: replace(__UpperCamelCase ,     __UpperCamelCase	) for k, v in initd.items()}
  assert _unmatched not in result.values(), "Incomplete partition spec."
  return freeze(unflatten_dict(__UpperCamelCase	)	)
