Metadata-Version: 2.1
Name: convlstm-mlx
Version: 0.1.1
Summary: An MLX implementation of ConvLSTM
Home-page: https://github.com/tomo-oga/convlstm-mlx
Author: Tomohiro Oga
Author-email: oga.t@northeastern.edu
License: MIT
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Development Status :: 3 - Alpha
Requires-Python: >=3.7
Description-Content-Type: text/markdown
Requires-Dist: mlx

[![PyPi Release](https://github.com/tomo-oga/convlstm-mlx/actions/workflows/python-publish.yml/badge.svg)](https://github.com/tomo-oga/convlstm-mlx/actions/workflows/python-publish.yml)

# An Implementation for ConvLSTM in Apple's Array Framework, MLX

A Convolutional LSTM recurrent layer. 


### `_conv_lstm_cell` 
This `nn.Module` computes the hidden and cell state for a time-step, expressed as:

$` i_t = \sigma (W_{xi} \ast X_t + W_{hi} \ast H_{t-1} + W_{ci} \odot C_{t-1} + b_i) `$ \
$` f_t = \sigma (W_{xf} \odot X_t \ast H_{t-1} + W_{cf} \odot C_{t-1} + b_f) `$\
$` C_t = f_t \odot X_t + i_t \odot tanh(W_{xc} \ast X_{t} + W_{hc} \ast H_{t-1} + b_c) `$ \
$` o_t = \sigma(W_{xo} \ast X_t +  W_{ho} \ast H_{t-1} + W_{co} \odot C_t + b_o )`$ \
$` H_t = o_t \odot tanh(C_t) `$ \

Where $`\sigma`$ and $`\odot`$ represent the hyperbolic sigmoid function and Hadamard product respectively.

The expected input for this layer has shape `NHWC` or  `HWC` where:

* `N` is the optional batch dimension
* `H` is the input's spatial height dimension
* `W` is the input's spatial width dimension
* `C` is the input's channel dimension

And returns a `Tuple` of the hidden state, $`H_t`$, and the cell state, $`C_t`$, each with shape `NHWO`. 

**Args:** 

`in_channels (int)`: The number of input channels, `C`.\
`out_channels (int)`: The number of output channels, `O`.\
`kernel_size (int)`: The size of the convolution filters, must be odd to keep spatial dimensions with padding. Default: `5`. \
`stride (Union[int, tuple]` : The stride of the convolution.
`padding (Union[int, tuple]` : Padding to add to the input for convolution.
`dilation (Union[int, tuple]` : Dilation of the convolution.
`bias` (bool): Whether the convolutional calculation should use biases or not. Default: `True.`

### `ConvLSTM`

Unrolls a `_conv_lstm_cell` sequentially over time-steps.

The expected input for this layer has shape `NLHWC` or  `LHWC` where:

* `N` is the optional batch dimension
* `L` is the length of the sequence
* `H` is the input's spatial height dimension
* `W` is the input's spatial width dimension
* `C` is the input's channel dimension

**Args:**
`in_channels (int)`: The number of input channels, `C`.\
`out_channels (int)`: The number of output channels, `O`.\
`kernel_size (int)`: The size of the convolution filters, must be odd to keep spatial dimensions with padding. Default: `5`. \
`bias` (bool): Whether the convolutional calculation should use biases or not. Default: `True.`

The following features are yet to be implemented from initial release:
- [ ] Bi-directionality - allows the conv-lstm to unroll both forwards and backwards across the sequence
- [ ] Allow for stride customization
- [ ] Allow for customizable padding along with modes 'same' and 'valid'




