# AUTOGENERATED FROM: babilim/model/layers/selection.ipynb

# Cell: 0
"""doc
# babilim.model.layers.selection

> These layers select parts of a tensor.
"""

# Cell: 1
from babilim.core.annotations import RunOnlyOnce
from babilim.core.module_native import ModuleNative


# Cell: 2
class Gather(ModuleNative):
    def __init__(self, axis):
        """
        Gather tensors from one tensor by providing an index tensor.
    
        Created object is callable with the following parameters:
        * **input_tensor**: (Tensor[N, L, ?]) The tensor from which to gather values at the given indices.
        * **indices**: (Tensor[N, K]) The indices at which to return the values of the input tensor.
        * **returns**: (Tensor[N, K, ?]) The tensor containing the values at the indices given.

        Arguments:
        :param axis: The axis along which to select.
        """
        super().__init__()
        assert axis != 0, "You cannot gather over the batch dimension."
        if axis > 0:
            axis = axis - 1
        self.axis = axis
        
    @RunOnlyOnce
    def _build_pytorch(self, input_tensor, indices):
        pass

    def _call_pytorch(self, input_tensor, indices):
        import torch
        assert input_tensor.shape[0] == indices.shape[0]

        # Then gather the indices along the batches.
        return torch.stack([torch.index_select(input_tensor[i], self.axis, indices[i]) for i in range(indices.shape[0])])
    
    @RunOnlyOnce
    def _build_tf(self, input_tensor, indices):
        # TODO implement
        raise NotImplementedError()
        
    def _call_tf(self, input_tensor, indices):
        # TODO implement
        raise NotImplementedError()


# Cell: 3
class TopKIndices(ModuleNative):
    def __init__(self, k):
        """
        Returns the top k tensor indices (separate per batch).
    
        Created object is callable with the following parameters:
        * **input_tensor**: (Tensor[N, L]) The tensor in which to search the top k indices.
        * **returns**: (Tensor[N, K]) The tensor containing the indices of the top k values.
        
        Parameters for the constructor:
        :param k: The number of indices to return per batch.
        """
        super().__init__()
        self.k = k
        
    @RunOnlyOnce
    def _build_pytorch(self, input_tensor):
        pass

    def _call_pytorch(self, input_tensor):
        import torch
        return torch.topk(input_tensor, self.k).indices

    @RunOnlyOnce
    def _build_tf(self, input_tensor):
        # TODO implement
        raise NotImplementedError()
        
    def _call_tf(self, input_tensor):
        # TODO implement
        raise NotImplementedError()
