def test_mu_law_companding(self):

        quantization_channels = 256

        waveform = self.waveform.clone()
        if not waveform.is_floating_point():
            waveform = waveform.to(torch.get_default_dtype())
        waveform /= torch.abs(waveform).max()

        self.assertTrue(waveform.min() >= -1. and waveform.max() <= 1.)

        waveform_mu = transforms.MuLawEncoding(quantization_channels)(waveform)
        self.assertTrue(waveform_mu.min() >= 0. and waveform_mu.max() <= quantization_channels)

        waveform_exp = transforms.MuLawDecoding(quantization_channels)(waveform_mu)
        self.assertTrue(waveform_exp.min() >= -1. and waveform_exp.max() <= 1.) 

def _test_istft_of_sine(self, amplitude, L, n):
        # stft of amplitude*sin(2*pi/L*n*x) with the hop length and window size equaling L
        x = torch.arange(2 * L + 1, dtype=torch.get_default_dtype())
        sound = amplitude * torch.sin(2 * math.pi / L * x * n)
        # stft = torch.stft(sound, L, hop_length=L, win_length=L,
        #                   window=torch.ones(L), center=False, normalized=False)
        stft = torch.zeros((L // 2 + 1, 2, 2))
        stft_largest_val = (amplitude * L) / 2.0
        if n < stft.size(0):
            stft[n, :, 1] = -stft_largest_val

        if 0 <= L - n < stft.size(0):
            # symmetric about L // 2
            stft[L - n, :, 1] = stft_largest_val

        estimate = torchaudio.functional.istft(stft, L, hop_length=L, win_length=L,
                                               window=torch.ones(L), center=False, normalized=False)
        # There is a larger error due to the scaling of amplitude
        _compare_estimate(sound, estimate, atol=1e-3) 

def __init__(self, manifold: Manifold, scale=1.0, learnable=False):
        super().__init__()
        self.base = manifold
        scale = torch.as_tensor(scale, dtype=torch.get_default_dtype())
        scale = scale.requires_grad_(False)
        if not learnable:
            self.register_buffer("_scale", scale)
            self.register_buffer("_log_scale", None)
        else:
            self.register_buffer("_scale", None)
            self.register_parameter("_log_scale", torch.nn.Parameter(scale.log()))
        # do not rebuild scaled functions very frequently, save them

        for method, scaling_info in self.base.__scaling__.items():
            # register rescaled functions as bound methods of this particular instance
            unbound_method = getattr(self.base, method).__func__  # unbound method
            self.__setattr__(
                method, types.MethodType(rescale(unbound_method, scaling_info), self)
            ) 

def penalty(self, s):
        """
        Calculate L1 Penalty.
        """
        to_return = torch.sum(s) / self.D
        if self.soft_groups is not None:
            # if soft_groups, there is an additional penalty for using more
            # groups
            s_grouped = torch.zeros(self.soft_D, 1,
                                    dtype=torch.get_default_dtype(),
                                    device=self.device)
            for group in torch.unique(self.soft_groups):
                # groups should be indexed 0 to n_group - 1
                # TODO: consider other functions here
                s_grouped[group] = s[self.soft_groups == group].max()
            # each component of the penalty contributes .5
            # TODO: could make this a user given parameter
            to_return = (to_return + torch.sum(s_grouped) / self.soft_D) * .5
        return to_return 

def forward_and_backward(self, s, xsub, ysub, retain_graph=False):
        """
        Completes the forward operation and computes gradients for learnability and penalty.
        """
        f_train = self.f_train(s, xsub, ysub)
        pen = self.penalty(s)
        # pylint: disable=E1102
        grad_outputs = torch.tensor([[1]], dtype=torch.get_default_dtype(),
                                    device=self.device)
        g1, = torch.autograd.grad([f_train], [self.x], grad_outputs,
                                  retain_graph=True)
        # pylint: disable=E1102
        grad_outputs = torch.tensor([[1]], dtype=torch.get_default_dtype(),
                                    device=self.device)
        g2, = torch.autograd.grad([pen], [self.x], grad_outputs,
                                  retain_graph=retain_graph)
        return f_train, pen, g1, g2 

def scale(self, waveform, factor=2.0**31):
        # scales a waveform by a factor
        if not waveform.is_floating_point():
            waveform = waveform.to(torch.get_default_dtype())
        return waveform / factor 

def use_floatX(request):
    dtype_old = torch.get_default_dtype()
    torch.set_default_dtype(request.param)
    yield request.param
    torch.set_default_dtype(dtype_old) 

def __enter__(self):
        """Set new dtype."""
        self.old_dtype = torch.get_default_dtype()
        torch.set_default_dtype(self.new_dtype) 

def test_set_torch_dtype():
    """Test dtype context manager."""
    dtype = torch.get_default_dtype()

    torch.set_default_dtype(torch.float32)
    with SetTorchDtype(torch.float64):
        a = torch.zeros(1)

    assert a.dtype is torch.float64
    b = torch.zeros(1)
    assert b.dtype is torch.float32

    torch.set_default_dtype(dtype) 

def __init__(self, *sizes, dtype=None, device=None):
        super(ZeroLazyTensor, self).__init__(*sizes)
        self.sizes = list(sizes)

        self._dtype = dtype or torch.get_default_dtype()
        self._device = device or torch.device("cpu") 

def dtype(self):
        if self.has_lengthscale:
            return self.lengthscale.dtype
        else:
            for param in self.parameters():
                return param.dtype
            return torch.get_default_dtype() 

def __init__(self, network, heading_dim=2, pre_norm=False, separate=True, get_prediction=False, weight=None):
        """
        Calculate heading angle with/out loss.
        The heading angle is absolute heading from the point person starts moving

        :param network: Network model
            - input - imu features
            - output - absolute_heading, prenorm_abs
        :param pre_norm: force network to output normalized values by adding a loss
        :param separate: report errors separately with total (The #losses can be obtained from get_channels()).
                        If False, only (weighted) sum of all losses will be returned.
        :param get_prediction: For testing phase, to get prediction values. In training only losses will be returned.
        :param weight: weight for each loss type (should ensure enough channels are given). If not all will be
                        weighed equally.
        """
        super(HeadingNetwork, self).__init__()
        self.network = network
        self.h_dim = heading_dim
        self.pre_norm = pre_norm
        self.concat_loss = not separate

        self.predict = get_prediction

        losses, channels = self.get_channels()
        if self.predict or weight is None or weight in ('None', 'none', False):
            self.weights = torch.ones(channels, dtype=torch.get_default_dtype(), device=_device)
        else:
            assert len(weight) == channels
            self.weights = torch.tensor(weight).to(_device) 

def get_init(data_train, init_type='on', rng=np.random.RandomState(0), prev_score=None):
    """
    Initialize the 'x' variable with different settings
    """

    D = data_train.n_features
    value_off = constants.Initialization.VALUE_DICT[
        constants.Initialization.OFF]
    value_on = constants.Initialization.VALUE_DICT[
        constants.Initialization.ON]

    if prev_score is not None:
        x0 = prev_score
    elif not isinstance(init_type, str):
        x0 = value_off * np.ones(D)
        x0[init_type] = value_on
    elif init_type.startswith(constants.Initialization.RANDOM):
        d = int(init_type.replace(constants.Initialization.RANDOM, ''))
        x0 = value_off * np.ones(D)
        x0[rng.permutation(D)[:d]] = value_on
    elif init_type == constants.Initialization.SKLEARN:
        B = data_train.return_raw
        X, y = data_train.get_dense_data()
        data_train.set_return_raw(B)
        ix = train_sk_dense(init_type, X, y, data_train.classification)
        x0 = value_off * np.ones(D)
        x0[ix] = value_on
    elif init_type in constants.Initialization.VALUE_DICT:
        x0 = constants.Initialization.VALUE_DICT[init_type] * np.ones(D)
    else:
        raise NotImplementedError(
            'init_type {0} not supported yet'.format(init_type))
    # pylint: disable=E1102
    return torch.tensor(x0.reshape((-1, 1)),
                        dtype=torch.get_default_dtype()) 

def __init__(self, Nminibatch, D, coeff, groups=None, binary=False,
                 device=constants.Device.CPU):
        super(LearnabilityMB, self).__init__()

        a = coeff / scipy.special.binom(Nminibatch, np.arange(coeff.size) + 2)
        self.order = a.size
        # pylint: disable=E1102
        self.a = torch.tensor(a, dtype=torch.get_default_dtype(), requires_grad=False)
        self.binary = binary

        self.a = self.a.to(device) 

def set_dense_X(self):
        if self.storage_level != constants.StorageLevel.DISK:
            if self.dense_size_gb <= self.MAXMEMGB:
                if self.storage_level == constants.StorageLevel.SPARSE:
                    self.X = self.X.toarray()
                self.X = torch.as_tensor(
                    self.X, dtype=torch.get_default_dtype())
                self.storage_level = constants.StorageLevel.DENSE 

def set_data_stats(self, Xmn, sv1, Xsd=1., ymn=0., ysd=1.):
        """
        Saves dataset stats to self to be used for preprocessing.
        """

        self.Xmn = torch.as_tensor(
            Xmn, dtype=torch.get_default_dtype()).to(self.device)
        self.sv1 = torch.as_tensor(
            sv1, dtype=torch.get_default_dtype()).to(self.device)
        self.Xsd = torch.as_tensor(
            Xsd, dtype=torch.get_default_dtype()).to(self.device)
        self.ymn = torch.as_tensor(
            ymn, dtype=torch.get_default_dtype()).to(self.device)
        self.ysd = torch.as_tensor(
            ysd, dtype=torch.get_default_dtype()).to(self.device) 

def __getitem__(self, idx):

        with torch.no_grad():
            X, y = self.getXy(idx)
            X = X.toarray() if scipy.sparse.issparse(X) else X

            X = torch.as_tensor(
                X, dtype=torch.get_default_dtype()).to(self.device)
            y = torch.as_tensor(
                y, dtype=torch.get_default_dtype()).to(self.device)

            if not self.return_raw:
                X, y = self.apply_preprocess(X, y)

            if self.classification and (
                    self.n_classes is None or self.n_classes == 2):
                y[y == 0] = -1

            if self.return_np:
                if constants.Device.CPU not in self.device:
                    X = X.cpu()
                    y = y.cpu()
                X = X.numpy()
                y = y.numpy()
                return X, y

            return X, y 

def __init__(self, name: str):
        """
        Args:
            name: the metric's name

        """
        super().__init__()
        self.name = name
        self._dtype = torch.get_default_dtype()
        self._device = torch.device('cpu') 

def step(self, closure):
        """Performs a single optimization step.

        Arguments:
            closure (callable): A closure that reevaluates the model
                and returns the loss.
        """
        assert len(self.param_groups) == 1

        def wrapped_closure(flat_params):
            """closure must call zero_grad() and backward()"""
            flat_params = torch.from_numpy(flat_params)
            flat_params = flat_params.to(torch.get_default_dtype())
            self._distribute_flat_params(flat_params)
            loss = closure()
            loss = loss.item()
            flat_grad = self._gather_flat_grad().cpu().detach().numpy()
            return loss, flat_grad.astype('float64')

        initial_params = self._gather_flat_params()
        initial_params = initial_params.cpu().detach().numpy()

        bounds = self._gather_flat_bounds()

        # Magic
        sol = sopt.minimize(wrapped_closure,
                            initial_params,
                            method='L-BFGS-B',
                            jac=True,
                            bounds=bounds)

        final_params = torch.from_numpy(sol.x)
        final_params = final_params.to(torch.get_default_dtype())
        self._distribute_flat_params(final_params) 

def __init__(self, dim, c=0.):
        super().__init__(1)
        self.register_buffer("dim", torch.as_tensor(dim, dtype=torch.int))
        self.register_buffer("c", torch.as_tensor(c, dtype=torch.get_default_dtype())) 

def meshgrid(*axes, batch=False):
    """
    See NumPy's or PyTorch's `meshgrid()`.

    :param axes: a list of N ints or torch vectors

    :return: a list of N :class:`Tensor`, of N dimensions each
    """

    device = None
    if not hasattr(axes, '__len__'):
        axes = [axes]
    if hasattr(axes[0], '__len__'):
        axes = axes[0]
    if hasattr(axes[0], 'device'):
        device = axes[0].device
    axes = list(axes)
    N = len(axes)
    for n in range(N):
        if not hasattr(axes[n], '__len__'):
            axes[n] = torch.arange(axes[n], dtype=torch.get_default_dtype())

    tensors = []
    for n in range(N):
        cores = [torch.ones(1, len(ax), 1).to(device) for ax in axes]
        if isinstance(axes[n], torch.Tensor):
            cores[n] = axes[n].type(torch.get_default_dtype())
        else:
            cores[n] = torch.tensor(axes[n].type(torch.get_default_dtype()))
        cores[n] = cores[n][None, :, None].to(device)
        tensors.append(tn.Tensor(cores, device=device, batch=batch))
    return tensors 

def arange(*args, **kwargs):
    """
    Creates a 1D :class:`Tensor` (see PyTorch's `arange`).

    :param args:
    :param kwargs:

    :return: a 1D :class:`Tensor`
    """

    return tn.Tensor([torch.arange(*args, dtype=torch.get_default_dtype(), **kwargs)[None, :, None]]) 

def __init__(self,
                 *array: Union[np.ndarray, pd.DataFrame, pd.Series,
                               torch.Tensor],
                 dtypes: Union[None, Sequence[torch.dtype]] = None):
        if dtypes is None:
            dtypes = [torch.get_default_dtype()] * len(array)
        if len(dtypes) != len(array):
            raise ValueError('length of dtypes not equal to length of array')

        array = [
            self._convert(data, dtype) for data, dtype in zip(array, dtypes)
        ]
        super().__init__(*array) 

def _test_backward(self, state, eps=2e-8, atol=1e-5, rtol=1e-3, max_num_per_param=5):
        @contextlib.contextmanager
        def double_prec():
            saved_dtype = torch.get_default_dtype()
            torch.set_default_dtype(torch.double)
            yield
            torch.set_default_dtype(saved_dtype)

        with double_prec():
            models = [m.to(torch.double) for m in networks.get_networks(state, 1)]
            trainer = Trainer(state, models)

            model = trainer.models[0]

            rdata, rlabel = next(iter(state.train_loader))
            rdata = rdata.to(state.device, torch.double, non_blocking=True)
            rlabel = rlabel.to(state.device, non_blocking=True)
            steps = trainer.get_steps()

            l, saved = trainer.forward(model, rdata, rlabel, steps)
            grad_info = trainer.backward(model, rdata, rlabel, steps, saved)
            trainer.accumulate_grad([grad_info])

            with torch.no_grad():
                for p_idx, p in enumerate(trainer.params):
                    pdata = p.data
                    N = p.numel()
                    for flat_i in np.random.choice(N, min(N, max_num_per_param), replace=False):
                        i = []
                        for s in reversed(p.size()):
                            i.insert(0, flat_i % s)
                            flat_i //= s
                        i = tuple(i)
                        ag = p.grad[i].item()
                        orig = pdata[i].item()
                        pdata[i] -= eps
                        steps = trainer.get_steps()
                        lm, _ = trainer.forward(model, rdata, rlabel, steps)
                        pdata[i] += eps * 2
                        steps = trainer.get_steps()
                        lp, _ = trainer.forward(model, rdata, rlabel, steps)
                        ng = (lp - lm).item() / (2 * eps)
                        pdata[i] = orig
                        rel_err = abs(ag - ng) / (atol + rtol * abs(ng))
                        info_msg = "testing param {} with shape [{}] at ({}):\trel_err={:.4f}\t" \
                                   "analytical={:+.6f}\tnumerical={:+.6f}".format(
                                       p_idx, format_intlist(p.size()),
                                       format_intlist(i), rel_err, ag, ng)
                        if unittest_verbosity() > 0:
                            print(info_msg)
                        self.assertTrue(rel_err <= 1, "gradcheck failed when " + info_msg) 

def totensor(array, dtype=None):
    """
    [[Source]](https://github.com/seba-1511/cherry/blob/master/cherry/_torch.py)

    **Description**

    Converts the argument `array` to a torch.tensor 1xN, regardless of its
    type or dimension.

    **Arguments**

    * **array** (int, float, ndarray, tensor) - Data to be converted to array.
    * **dtype** (dtype, *optional*, default=None) - Data type to use for representation.
    By default, uses `torch.get_default_dtype()`.

    **Returns**

    * Tensor of shape 1xN with the appropriate data type.

    **Example**

    ~~~python
    array = [5, 6, 7.0]
    tensor = cherry.totensor(array, dtype=th.float32)
    array = np.array(array, dtype=np.float64)
    tensor = cherry.totensor(array, dtype=th.float16)
    ~~~

    """
    if dtype is None:
        dtype = th.get_default_dtype()
    if isinstance(array, (list, tuple)):
        array = th.cat([totensor(x) for x in array], dim=0)
    if isinstance(array, int):
        array = float(array)
    if isinstance(array, float):
        array = [array, ]
    if isinstance(array, list):
        array = np.array(array)
    if isinstance(array, (np.ndarray,
                          np.bool_,
                          np.float32,
                          np.float64,
                          np.int32,
                          np.int64)):
        if array.dtype == np.bool_:
            array = array.astype(np.uint8)
        array = th.tensor(array, dtype=dtype)
        array = array.unsqueeze(0)
    while len(array.shape) < 2:
        array = array.unsqueeze(0)
    return array 

def _partial_fit(self, X, y, n_classes=None, groups=None):
        """
        Private function for partial_fit to enable trying floats before doubles.
        """
        # pass in X and y in chunks
        if hasattr(self, '_data_train'):
            # just overwrite the X and y from the new chunk but make them tensors
            # keep dataset stats from previous
            self._data_train.X = X.values if isinstance(X, pd.DataFrame) else X
            self._data_train.N, self._data_train.D = self._data_train.X.shape
            self._data_train.dense_size_gb = self._data_train.get_dense_size()
            self._data_train.set_dense_X()

            self._data_train.y = y.values if isinstance(y, pd.Series) else y
            self._data_train.y = torch.as_tensor(
                y, dtype=torch.get_default_dtype())
        else:
            data_train = self._prepare_data(X, y, n_classes=n_classes)
            self._data_train = data_train

        batch_size, _, accum_steps, max_iter = self._set_batch_size(
            self._data_train)

        rng = None  # not used
        debug = 0  # {0,1} print messages and do other stuff?
        dn_logs = None  # tensorboard logs; only specify if debug=1
        path_save = None  # intermediate models saves; only specify if debug=1
        m, solver = _train(self._data_train,
                           batch_size,
                           self.order,
                           self.penalty,
                           rng,
                           self.learning_rate,
                           debug,
                           max_iter,
                           self.max_time,
                           self.init,
                           self.dftol_stop,
                           self.freltol_stop,
                           dn_logs,
                           accum_steps,
                           path_save,
                           self.shuffle,
                           device=self.device,
                           verbose=self.verbose,
                           prev_checkpoint=self._prev_checkpoint if hasattr(
                               self, '_prev_checkpoint') else None,
                           groups=groups if not self.soft_grouping else None,
                           soft_groups=groups if self.soft_grouping else None)

        self._prev_checkpoint = m
        self._process_results(m, solver, X, groups=groups)
        return self 

def compute_data_stats(self):
        """
        1. computes/estimates feature means
        2. if preprocess == 'zscore', computes/estimates feature standard devs
        3. if not classification, computes/estimates target mean/standard dev
        4. estimates largest singular value of data matrix
        """
        t = time.time()
        X, y = self.X[self.ix_statistics], self.y[self.ix_statistics]
        preprocess = self.preprocess
        classification = self.classification

        Xmn = (X.mean(dim=0)
               if not scipy.sparse.issparse(X)
               else np.array(X.mean(axis=0)).ravel())

        if preprocess == constants.Preprocess.ZSCORE:
            Xsd = (X.std(dim=0)
                   if not scipy.sparse.issparse(X)
                   else PrepareData.sparse_std(X, Xmn))
            Xsd[Xsd == 0] = 1.
        else:
            Xsd = 1.

        if preprocess is not None and preprocess:
            if preprocess == constants.Preprocess.ZSCORE:
                Xc = (X - Xmn) / Xsd
            else:
                Xc = X - Xmn
        else:
            Xc = X - Xmn

        sv1 = scipy.sparse.linalg.svds(Xc / (
            torch.sqrt(torch.prod(torch.as_tensor(y.size(), dtype=torch.get_default_dtype())))
            if not scipy.sparse.issparse(X) else y.numpy().size),
                                       k=1,
                                       which='LM',
                                       return_singular_vectors=False)
        # avoid runaway sv1
        sv1 = np.array([min(np.finfo(np.float32).max,
                            sv1[0])])

        if not classification:
            ymn = y.mean()
            ysd = y.std()
        else:
            # TODO: set these, for each class?
            ymn = 0.
            ysd = 1.
        if self.verbose:
            print(" computing data statistics took: ", time.time() - t)

        return Xmn, sv1, Xsd, ymn, ysd 

def _full_rank_tt(data, batch=False): # Naive TT formatting, don't even attempt to compress
    data = data.to(torch.get_default_dtype())
    shape = data.shape
    result = []

    if batch:
        N = data.dim() - 1
    else:
        N = data.dim()

    data = data if isinstance(data, torch.Tensor) else torch.tensor(data)
    device = data.device

    if batch:
        resh = torch.reshape(data, [shape[0], shape[1], -1])
    else:
        resh = torch.reshape(data, [shape[0], -1])

    for n in range(1, N):
        if batch:
            if resh.shape[1] < resh.shape[2]:
                I = torch.cat([torch.eye(resh.shape[1])[None, ...] for _ in range(resh.shape[0])]).to(device)
                result.append(torch.reshape(I, [resh.shape[0], resh.shape[1] // shape[n], shape[n], resh.shape[1]]))
                resh = torch.reshape(resh, (resh.shape[0], resh.shape[1] * shape[n + 1], resh.shape[2] // shape[n + 1]))
            else:
                result.append(torch.reshape(resh, [resh.shape[0], resh.shape[1] // shape[n], shape[n], resh.shape[2]]))
                I = torch.cat([torch.eye(resh.shape[2])[None, ...] for _ in range(resh.shape[0])]).to(device)
                resh = torch.reshape(I, (resh.shape[0], resh.shape[2] * shape[n + 1], resh.shape[2] // shape[n + 1]))
        else:
            if resh.shape[0] < resh.shape[1]:
                result.append(torch.reshape(torch.eye(resh.shape[0]).to(device), [resh.shape[0] // shape[n - 1],
                                                                           shape[n - 1], resh.shape[0]]))
                resh = torch.reshape(resh, (resh.shape[0] * shape[n], resh.shape[1] // shape[n]))
            else:
                result.append(torch.reshape(resh, [resh.shape[0] // shape[n - 1],
                                                       shape[n - 1], resh.shape[1]]))
                resh = torch.reshape(torch.eye(resh.shape[1]).to(device), (resh.shape[1] * shape[n], resh.shape[1] // shape[n]))

    if batch:
        result.append(torch.reshape(resh, [resh.shape[0], resh.shape[1] // shape[N], shape[N], 1]))
    else:
        result.append(torch.reshape(resh, [resh.shape[0] // shape[N - 1], shape[N - 1], 1]))
    return result 

def __init__(
        self,
        *,
        x_dtype: Union[torch.dtype, Sequence[torch.dtype]] = None,
        y_dtype: Union[torch.dtype, Sequence[torch.dtype]] = None,
        empty_cache: bool = False,
        auto_reshape: bool = True,
        argmax: bool = False,
    ):
        """
        Covert tensor like data into :class:`torch.Tensor` automatically.
        
        Parameters
        ----------
        x_dtype
            The :class:`torch.dtype`s of **X** data.
            If ``None``, will convert all data into ``torch.get_default_dtype()`` type.
            Can be a tuple of ``torch.dtype`` when your **X** is a tuple.
        y_dtype
            The :class:`torch.dtype`s of **y** data.
            If ``None``, will convert all data into ``torch.get_default_dtype()`` type.
            Can be a tuple of ``torch.dtype`` when your **y** is a tuple.
        empty_cache
            See Also: https://pytorch.org/docs/stable/cuda.html#torch.cuda.empty_cache
        auto_reshape
            Reshape tensor to (-1, 1) if tensor shape is (n,). Default ``True``.
        argmax
            Apply ``np.argmax(out, 1)`` on the output. This should only be used with classification model.
        """

        self._argmax = argmax
        self._auto_reshape = False if argmax else auto_reshape
        self._empty_cache = empty_cache
        if x_dtype is None:
            self._x_dtype = torch.get_default_dtype()
        else:
            self._x_dtype = x_dtype

        if y_dtype is None:
            self._y_dtype = torch.get_default_dtype()
        else:
            self._y_dtype = y_dtype