# 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 axis != -1:
    raise NotImplementedError(
        'Only last axis supported for now. Please send a feature request.')

  # Normalize to array (to support broadcasting).
  # If inputs are static arguments (not `tf.Tensor`), use numpy arrays for
  # optimization (factors statically computed).
  from_ = tuple(_to_array(v) for v in from_)
  to = tuple(_to_array(v) for v in to)

  # `a` can be scalar or array of shape=(x.shape[-1],), same for `b`
  a, b = _linear_interp_factors(*from_, *to)
  return a * x + b


def _linear_interp_factors(
    old_min: _MinMaxValue,
    old_max: _MinMaxValue,
    new_min: _MinMaxValue,
    new_max: _MinMaxValue,
) -> Tuple[Union[float, np.ndarray], Union[float, np.ndarray]]:
  """Resolve the `y = a * x + b` equation and returns the factors."""
  a = (new_min - new_max) / (old_min - old_max)
  b = (old_min * new_max - new_min * old_max) / (old_min - old_max)
  return a, b


def random_choice(
    a: Union[int, TensorLike],
    *,

# 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 axis != -1:
    raise NotImplementedError(
        'Only last axis supported for now. Please send a feature request.')

  # Normalize to array (to support broadcasting).
  # If inputs are static arguments (not `tf.Tensor`), use numpy arrays for
  # optimization (factors statically computed).
  from_ = tuple(_to_array(v) for v in from_)
  to = tuple(_to_array(v) for v in to)

  # `a` can be scalar or array of shape=(x.shape[-1],), same for `b`
  a, b = _linear_interp_factors(*from_, *to)
  return a * x + b


def _linear_interp_factors(
    old_min: _MinMaxValue,
    old_max: _MinMaxValue,
    new_min: _MinMaxValue,
    new_max: _MinMaxValue,
) -> Tuple[Union[float, np.ndarray], Union[float, np.ndarray]]:
  """Resolve the `y = a * x + b` equation and returns the factors."""
  a = (new_min - new_max) / (old_min - old_max)
  b = (old_min * new_max - new_min * old_max) / (old_min - old_max)
  return a, b


def random_choice(
    a: Union[int, TensorLike],
    *,

    size: Union[None, int, Shape] = None,