# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04_utils.ipynb (unless otherwise specified).

__all__ = ['to_device', 'device', 'to_np', 'init_net']

# Cell
import torch
from torch import Tensor
from torch.nn import init
import torch.nn as nn

from typing import Union, List
import numpy as np

# Cell
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

def to_device(x:Union[Tensor,List[Tensor]], device:torch.device=device) -> Union[Tensor,List[Tensor]]:
    r'''
    Recursively place Tensor(s) onto device

    Arguments:
        x: Tensor(s) to place on device

    Returns:
        Tensor(s) on device
    '''

    if x is None: return x
    if isinstance(x, list): return [to_device(o, device) for o in x]
    return x.to(device)

# Cell
def to_np(x:Tensor) -> np.ndarray:
    r'''
    Convert Tensor x to a Numpy array

    Arguments:
        x: Tensor to convert

    Returns:
        x as a Numpy array
    '''

    return x.cpu().detach().numpy()

# Cell
def init_net(model:nn.Module):
    r'''Recursively initialise fully-connected ReLU network with Kaiming and zero bias'''
    if isinstance(model,nn.Linear):
        init.kaiming_normal_(model.weight, nonlinearity='relu')
        init.zeros_(model.bias)
    for l in model.children(): init_net(l)