# Path: sunds/typing.py



import functools
import operator
import numpy as np
import tensorflow as tf
from typing import Sequence, Tuple, Union
from sunds.typing import Shape, Tensor, TensorLike  # pylint: disable=g-multiple-import
  """
  if isinstance(a, int):  # Normalize int
    a = tf.range(a)
  a = tf.convert_to_tensor(a)
  if replace:
    raise NotImplementedError(
        'Only replace=False supported. You can use `tf.random.categorical.`')
  if len(a.shape) != 1:
    raise NotImplementedError(
        '`random_choice` only support single dim tensors.')

  shape = (size,) if isinstance(size, int) else tuple(size)
  # TODO(py3.8): Replace by `math.prod`
  num_samples = functools.reduce(operator.mul, shape)
  num_values = a.shape[0]  # Numbe of values to sample
  if not replace and num_samples > num_values:
    raise ValueError(
        'Cannot take a larger sample than population when `replace=False`. '
        f'{num_samples} > {num_values}')

  indices = tf.range(num_values)
  indices = tf.random.shuffle(indices)
  indices = indices[:num_samples]

  a = tf.gather(a, indices)
  return tf.reshape(a, shape)


def _to_array(x: _MinMaxValue) -> _MinMaxValue:
  """Convert to array/tensor."""

# Path: sunds/typing.py



# Path: sunds/utils/tf_utils.py
import functools
import operator
import numpy as np
import tensorflow as tf
from typing import Sequence, Tuple, Union
from sunds.typing import Shape, Tensor, TensorLike  # pylint: disable=g-multiple-import

  """
  if isinstance(a, int):  # Normalize int
    a = tf.range(a)
  a = tf.convert_to_tensor(a)
  if replace:
    raise NotImplementedError(
        'Only replace=False supported. You can use `tf.random.categorical.`')
  if len(a.shape) != 1:
    raise NotImplementedError(
        '`random_choice` only support single dim tensors.')

  shape = (size,) if isinstance(size, int) else tuple(size)
  # TODO(py3.8): Replace by `math.prod`
  num_samples = functools.reduce(operator.mul, shape)
  num_values = a.shape[0]  # Numbe of values to sample
  if not replace and num_samples > num_values:
    raise ValueError(
        'Cannot take a larger sample than population when `replace=False`. '
        f'{num_samples} > {num_values}')

  indices = tf.range(num_values)
  indices = tf.random.shuffle(indices)
  indices = indices[:num_samples]

  a = tf.gather(a, indices)
  return tf.reshape(a, shape)


def _to_array(x: _MinMaxValue) -> _MinMaxValue:
  """Convert to array/tensor."""

  if isinstance(x, tf.Tensor):