# AUTOGENERATED FROM: babilim/core/tensor.ipynb

# Cell: 0
"""doc
# babilim.core.tensor

> Create a tensor independent of the underlying framework.
"""

# Cell: 1
from typing import Union, Any
import numpy as np
import babilim
from babilim import PYTORCH_BACKEND, TF_BACKEND
from babilim.core.itensor import ITensor, ITensorWrapper


# Cell: 2
"""doc
# Create Tensors

This package creates tensors of type ITensor.
It does not contain any classes but just a function which creates tensors.
This is because there is different implementations of the ITensor interface for pytorch and tensorflow required.
"""

# Cell: 3
def Tensor(*, data: Union[np.ndarray, Any], trainable: bool) -> ITensor:
    """
    Create a babilim tensor from a native tensor or numpy array.

    :param data: The data that should be put in a babilim tensor. This can be either a numpy array or a pytorch/tensorflow tensor.
    :param trainable: If the tensor created should be trainable. Only works for numpy tensors, native tensors overwrite this field!
    :return: An object of type babilim.core.ITensor.
    """
    if babilim.get_backend() == PYTORCH_BACKEND:
        from babilim.core.tensor_pt import Tensor as _Tensor
        from torch import Tensor as _PtTensor
        native = None
        if isinstance(data, _PtTensor):
            native = data
            data = None
        return _Tensor(data, trainable, native)
    elif babilim.get_backend() == TF_BACKEND:
        from babilim.core.tensor_tf import Tensor as _Tensor
        from tensorflow import Tensor as _TfTensor
        native = None
        if isinstance(data, _TfTensor):
            native = data
            data = None
        return _Tensor(data, trainable, native)
    else:
        raise RuntimeError("No variable implementation for this backend was found. (backend={})".format(babilim.get_backend()))


# Cell: 4
def TensorWrapper() -> ITensorWrapper:
    """
    Create a tensor wrapper object.
    
    Sometimes it is nescesarry to implement stuff in native pytorch or native tensorflow. Here the tensor wrapper can help.

    **WARNING: Instead of directly using the TensorWrapper, you should prefer using the babilim.module.Lambda!**
    """
    if babilim.get_backend() == PYTORCH_BACKEND:
        from babilim.core.tensor_pt import TensorWrapper as _TensorWrapper
        return _TensorWrapper()
    elif babilim.get_backend() == TF_BACKEND:
        from babilim.core.tensor_tf import TensorWrapper as _TensorWrapper
        return _TensorWrapper()
    else:
        raise RuntimeError("No variable implementation for this backend was found. (backend={})".format(babilim.get_backend()))
