.. raw:: html

    <p align="center">
        <img width="75%" style="display: block;" src='docs/partial_source/logos/logo.png'>
    </p>

.. raw:: html

    <br/>
    <a href="https://pypi.org/project/ivy-memory">
        <img style="float: left; padding-right: 4px; padding-bottom: 4px;" src="https://badge.fury.io/py/ivy-memory.svg">
    </a>
    <a href="https://github.com/ivy-dl/memory/actions?query=workflow%3Adocs">
        <img style="float: left; padding-right: 4px; padding-bottom: 4px;" src="https://img.shields.io/github/workflow/status/ivy-dl/memory/docs?label=docs">
    </a>
    <a href="https://github.com/ivy-dl/memory/actions?query=workflow%3Anightly-tests">
        <img style="float: left; padding-right: 4px; padding-bottom: 4px;" src="https://img.shields.io/github/workflow/status/ivy-dl/memory/nightly-tests?label=tests">
    </a>
    <a href="https://discord.gg/EN9YS3QW8w">
        <img style="float: left; padding-right: 4px; padding-bottom: 4px;" src="https://img.shields.io/discord/799879767196958751?color=blue&label=%20&logo=discord&logoColor=white">
    </a>
    <br clear="all" />

**End-to-end memory modules for deep learning developers, written in Ivy.**

.. raw:: html

    <div style="display: block;">
        <img width="4%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/empty.png">
        <a href="https://jax.readthedocs.io">
            <img width="12%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/jax_logo.png">
        </a>
        <img width="6.5%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/empty.png">
        <a href="https://www.tensorflow.org">
            <img width="12%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/tensorflow_logo.png">
        </a>
        <img width="6.5%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/empty.png">
        <a href="https://pytorch.org">
            <img width="12%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/pytorch_logo.png">
        </a>
        <img width="6.5%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/empty.png">
        <a href="https://mxnet.apache.org">
            <img width="12%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/mxnet_logo.png">
        </a>
        <img width="6.5%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/empty.png">
        <a href="https://numpy.org">
            <img width="12%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/supported/numpy_logo.png">
        </a>
    </div>

Contents
--------

* `Overview`_
* `Run Through`_
* `Interactive Demos`_
* `Get Involed`_

Overview
--------

.. _docs: https://ivy-dl.org/memory

**What is Ivy Memory?**

Ivy memory provides differentiable memory modules,
including learnt modules such as Neural Turing Machines (NTM),
but also parameter-free modules such as End-to-End Egospheric Spatial Memory (ESM).
Check out the docs_ for more info!

The library is built on top of the Ivy deep learning framework.
This means all memory modules simultaneously support:
Jax, Tensorflow, PyTorch, MXNet, and Numpy.

**Ivy Libraries**

There are a host of derived libraries written in Ivy, in the areas of mechanics, 3D vision, robotics,
differentiable memory, and differentiable gym environments. Click on the icons below for their respective github pages.

.. raw:: html

    <div style="display: block;">
        <img width="9%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/mech">
            <img width="15%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/ivy_mech.png">
        </a>
        <img width="7%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/vision">
            <img width="15%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/ivy_vision.png">
        </a>
        <img width="7%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/robot">
            <img width="15%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/ivy_robot.png">
        </a>
        <img width="7%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/gym">
            <img width="15%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/ivy_gym.png">
        </a>

        <br clear="all" />

        <img width="10%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://pypi.org/project/ivy-mech">
            <img width="13%" style="float: left;" src="https://badge.fury.io/py/ivy-mech.svg">
        </a>
        <img width="9%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://pypi.org/project/ivy-vision">
            <img width="13%" style="float: left;" src="https://badge.fury.io/py/ivy-vision.svg">
        </a>
        <img width="9%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://pypi.org/project/ivy-robot">
            <img width="13%" style="float: left;" src="https://badge.fury.io/py/ivy-robot.svg">
        </a>
        <img width="9%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://pypi.org/project/ivy-gym">
            <img width="13%" style="float: left;" src="https://badge.fury.io/py/ivy-gym.svg">
        </a>

        <br clear="all" />

        <img width="12%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/mech/actions?query=workflow%3Anightly-tests">
            <img width="9%" style="float: left;" src="https://img.shields.io/github/workflow/status/ivy-dl/mech/nightly-tests?label=tests">
        </a>
        <img width="13%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/vision/actions?query=workflow%3Anightly-tests">
            <img width="9%" style="float: left;" src="https://img.shields.io/github/workflow/status/ivy-dl/vision/nightly-tests?label=tests">
        </a>
        <img width="13%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/robot/actions?query=workflow%3Anightly-tests">
            <img width="9%" style="float: left;" src="https://img.shields.io/github/workflow/status/ivy-dl/robot/nightly-tests?label=tests">
        </a>
        <img width="13%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/gym/actions?query=workflow%3Anightly-tests">
            <img width="9%" style="float: left;" src="https://img.shields.io/github/workflow/status/ivy-dl/gym/nightly-tests?label=tests">
        </a>

        <br clear="all" />

        <img width="20%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/memory">
            <img width="15%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/ivy_memory.png">
        </a>
        <img width="7%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/builder">
            <img width="15%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/ivy_builder.png">
        </a>
        <img width="7%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/models">
            <img width="15%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/ivy_models.png">
        </a>

        <br clear="all" />

        <img width="21%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://pypi.org/project/ivy-memory">
            <img width="13%" style="float: left;" src="https://badge.fury.io/py/ivy-memory.svg">
        </a>
        <img width="9%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://pypi.org/project/ivy-builder">
            <img width="13%" style="float: left;" src="https://badge.fury.io/py/ivy-builder.svg">
        </a>
        <img width="9%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://pypi.org/project/ivy-models">
            <img width="13%" style="float: left;" src="https://badge.fury.io/py/ivy-models.svg">
        </a>

        <br clear="all" />

        <img width="23%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/memory/actions?query=workflow%3Anightly-tests">
            <img width="9%" style="float: left;" src="https://img.shields.io/github/workflow/status/ivy-dl/memory/nightly-tests?label=tests">
        </a>
        <img width="13%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/builder/actions?query=workflow%3Anightly-tests">
            <img width="9%" style="float: left;" src="https://img.shields.io/github/workflow/status/ivy-dl/builder/nightly-tests?label=tests">
        </a>
        <img width="13%" style="float: left;" src="https://raw.githubusercontent.com/ivy-dl/ivy-dl.github.io/master/img/externally_linked/logos/empty.png">
        <a href="https://github.com/ivy-dl/models/actions?query=workflow%3Anightly-tests">
            <img width="9%" style="float: left;" src="https://img.shields.io/github/workflow/status/ivy-dl/models/nightly-tests?label=tests">
        </a>

        <br clear="all" />

    </div>
    <br clear="all" />

**Quick Start**

Ivy memory can be installed like so: ``pip install ivy-memory``

.. _demos: https://github.com/ivy-dl/memory/tree/master/ivy_memory_demos
.. _interactive: https://github.com/ivy-dl/memory/tree/master/ivy_memory_demos/interactive

To quickly see the different aspects of the library, we suggest you check out the demos_!
We suggest you start by running the script ``run_through.py``,
and read the "Run Through" section below which explains this script.

For more interactive demos, we suggest you run either
``learning_to_copy_with_ntm.py`` or ``mapping_a_room_with_esm.py`` in the interactive_ demos folder.

Run Through
-----------

We run through some of the different parts of the library via a simple ongoing example script.
The full script is available in the demos_ folder, as file ``run_through.py``.

**End-to-End Egospheric Spatial Memory**

First, we show how the Ivy End-to-End Egospheric Spatial Memory (ESM) class can be used inside a pure-Ivy model.
We first define the model as below.

.. code-block:: python

    class IvyModelWithESM(ivy.Module):

        def __init__(self, channels_in, channels_out):
            self._channels_in = channels_in
            self._esm = ivy_mem.ESM(omni_image_dims=(16, 32))
            self._linear = ivy_mem.Linear(channels_in, channels_out)
            ivy.Module.__init__(self, 'cpu')

        def _forward(self, obs):
            mem = self._esm(obs)
            x = ivy.reshape(mem.mean, (-1, self._channels_in))
            return self._linear(x)

Next, we instantiate this model, and verify that the returned tensors are of the expected shape.

.. code-block:: python

    # create model
    in_channels = 32
    out_channels = 8
    ivy.set_framework('torch')
    model = IvyModelWithESM(in_channels, out_channels)

    # input config
    batch_size = 1
    image_dims = [5, 5]
    num_timesteps = 2
    num_feature_channels = 3

    # create image of pixel co-ordinates
    uniform_pixel_coords =\
        ivy_vision.create_uniform_pixel_coords_image(image_dims, [batch_size, num_timesteps])

    # define camera measurement
    depths = ivy.random_uniform(shape=[batch_size, num_timesteps] + image_dims + [1])
    pixel_coords = ivy_vision.depth_to_pixel_coords(depths)
    inv_calib_mats = ivy.random_uniform(shape=[batch_size, num_timesteps, 3, 3])
    cam_coords = ivy_vision.pixel_to_cam_coords(pixel_coords, inv_calib_mats)[..., 0:3]
    features = ivy.random_uniform(shape=[batch_size, num_timesteps] + image_dims + [num_feature_channels])
    img_mean = ivy.concatenate((cam_coords, features), -1)
    cam_rel_mat = ivy.identity(4, batch_shape=[batch_size, num_timesteps])[..., 0:3, :]

    # place these into an ESM camera measurement container
    esm_cam_meas = ESMCamMeasurement(
        img_mean=img_mean,
        cam_rel_mat=cam_rel_mat
    )

    # define agent pose transformation
    agent_rel_mat = ivy.identity(4, batch_shape=[batch_size, num_timesteps])[..., 0:3, :]

    # collect together into an ESM observation container
    esm_obs = ESMObservation(
        img_meas={'camera_0': esm_cam_meas},
        agent_rel_mat=agent_rel_mat
    )

    # call model and test output
    output = model(esm_obs)
    assert output.shape[-1] == out_channels

Finally, we define a dummy loss function, and show how the ESM network can be trained using Ivy functions only.

.. code-block:: python

    # define loss function
    target = ivy.zeros_like(output)

    def loss_fn(v):
        pred = model(esm_obs, v=v)
        return ivy.reduce_mean((pred - target) ** 2)

    # optimizer
    optimizer = ivy.SGD(lr=1e-4)

    # train model
    print('\ntraining dummy Ivy ESM model...\n')
    for i in range(10):
        loss, grads = ivy.execute_with_gradients(loss_fn, model.v)
        model.v = optimizer.step(model.v, grads)
        print('step {}, loss = {}'.format(i, ivy.to_numpy(loss).item()))
    print('\ndummy Ivy ESM model trained!\n')
    ivy.unset_framework()

**Neural Turing Machine**

We next show how the Ivy Neural Turing Machine (NTM) class can be used inside a TensorFlow model.
First, we define the model as below.

.. code-block:: python

    class TfModelWithNTM(tf.keras.Model):

        def __init__(self, channels_in, channels_out):
            tf.keras.Model.__init__(self)
            self._linear = tf.keras.layers.Dense(64)
            memory_size = 4
            memory_vector_dim = 1
            self._ntm = ivy_mem.NTM(
                input_dim=64, output_dim=channels_out, ctrl_output_size=channels_out, ctrl_layers=1,
                memory_size=memory_size, memory_vector_dim=memory_vector_dim, read_head_num=1, write_head_num=1)
            self._assign_variables()

        def _assign_variables(self):
            self._ntm.v.map(
                lambda x, kc: self.add_weight(name=kc, shape=x.shape))
            self.set_weights([ivy.to_numpy(v) for k, v in self._ntm.v.to_iterator()])
            self.trainable_weights_dict = dict()
            for weight in self.trainable_weights:
                self.trainable_weights_dict[weight.name] = weight
            self._ntm.v = self._ntm.v.map(lambda x, kc: self.trainable_weights_dict[kc + ':0'])

        def call(self, x, **kwargs):
            x = self._linear(x)
            return self._ntm(x)

Next, we instantiate this model, and verify that the returned tensors are of the expected shape.

.. code-block:: python

    # create model
    in_channels = 32
    out_channels = 8
    ivy.set_framework('tensorflow')
    model = TfModelWithNTM(in_channels, out_channels)

    # define inputs
    batch_shape = [1, 2]
    timesteps = 3
    input_shape = batch_shape + [timesteps, in_channels]
    input_seq = tf.random.uniform(batch_shape + [timesteps, in_channels])

    # call model and test output
    output_seq = model(input_seq)
    assert input_seq.shape[:-1] == output_seq.shape[:-1]
    assert input_seq.shape[-1] == in_channels
    assert output_seq.shape[-1] == out_channels

Finally, we define a dummy loss function, and show how the NTM can be trained using a native TensorFlow optimizer.

.. code-block:: python

    # define loss function
    target = tf.zeros_like(output_seq)

    def loss_fn():
        pred = model(input_seq)
        return tf.reduce_sum((pred - target) ** 2)

    # define optimizer
    optimizer = tf.keras.optimizers.Adam(1e-2)

    # train model
    print('\ntraining dummy TensorFlow NTM model...\n')
    for i in range(10):
        with tf.GradientTape() as tape:
            loss = loss_fn()
        grads = tape.gradient(loss, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))
        print('step {}, loss = {}'.format(i, loss))
    print('\ndummy TensorFlow NTM model trained!\n')
    ivy.unset_framework()

Interactive Demos
-----------------

In addition to the run through above, we provide two further demo scripts,
which are more visual and interactive.

The scripts for these demos can be found in the interactive_ demos folder.

**Learning to Copy with NTM**

The first demo trains a Neural Turing Machine to copy a sequence from one memory bank to another.
NTM can overfit to a single copy sequence very quickly, as show in the real-time visualization below.

.. raw:: html

    <p align="center">
        <img width="75%" style="display: block;" src='https://github.com/ivy-dl/ivy-dl.github.io/blob/master/img/externally_linked/ivy_memory/demo_a.gif?raw=true'>
    </p>

**Mapping a Room with ESM**

The second demo creates an egocentric map of a room, from a rotating camera.
The raw image observations are shown on the left,
and the incrementally constructed omni-directional ESM feature and depth images are shown on the right.
While this example only projects color values into the memory, arbitrary neural network features can also be projected, for end-to-end training.

.. raw:: html

    <p align="center">
        <img width="75%" style="display: block;" src='https://github.com/ivy-dl/ivy-dl.github.io/blob/master/img/externally_linked/ivy_memory/demo_b.gif?raw=true'>
    </p>

Get Involed
-----------

We hope the memory classes in this library are useful to a wide range of deep learning developers.
However, there are many more areas of differentiable memory which could be covered by this library.

If there are any particular functions or classes you feel are missing,
and your needs are not met by the functions currently on offer,
then we are very happy to accept pull requests!

We look forward to working with the community on expanding and improving the Ivy memory library.

Citation
--------

::

    @article{lenton2021ivy,
      title={Ivy: Templated Deep Learning for Inter-Framework Portability},
      author={Lenton, Daniel and Pardo, Fabio and Falck, Fabian and James, Stephen and Clark, Ronald},
      journal={arXiv preprint arXiv:2102.02886},
      year={2021}
    }
