{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# For tips on running notebooks in Google Colab, see\n",
    "# https://pytorch.org/tutorials/beginner/colab\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(Prototype) MaskedTensor Sparsity\n",
    "=================================\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before working on this tutorial, please make sure to review our\n",
    "[MaskedTensor Overview tutorial\n",
    "\\<https://pytorch.org/tutorials/prototype/maskedtensor\\_overview.html\\>]{.title-ref}.\n",
    "\n",
    "Introduction\n",
    "============\n",
    "\n",
    "Sparsity has been an area of rapid growth and importance within PyTorch;\n",
    "if any sparsity terms are confusing below, please refer to the [sparsity\n",
    "tutorial](https://pytorch.org/docs/stable/sparse.html) for additional\n",
    "details.\n",
    "\n",
    "Sparse storage formats have been proven to be powerful in a variety of\n",
    "ways. As a primer, the first use case most practitioners think about is\n",
    "when the majority of elements are equal to zero (a high degree of\n",
    "sparsity), but even in cases of lower sparsity, certain formats (e.g.\n",
    "BSR) can take advantage of substructures within a matrix.\n",
    "\n",
    "<div style=\"background-color: #54c7ec; color: #fff; font-weight: 700; padding-left: 10px; padding-top: 5px; padding-bottom: 5px\"><strong>NOTE:</strong></div>\n",
    "\n",
    "<div style=\"background-color: #f3f4f7; padding-left: 10px; padding-top: 10px; padding-bottom: 10px; padding-right: 10px\">\n",
    "\n",
    "<p>At the moment, MaskedTensor supports COO and CSR tensors with plans to support additional formats(such as BSR and CSC) in the future. If you have any requests for additional formats,please file a feature request <a href=\"https://github.com/pytorch/pytorch/issues\">here</a>!</p>\n",
    "\n",
    "</div>\n",
    "\n",
    "Principles\n",
    "==========\n",
    "\n",
    "When creating a `MaskedTensor`{.interpreted-text role=\"class\"} with\n",
    "sparse tensors, there are a few principles that must be observed:\n",
    "\n",
    "1.  `data` and `mask` must have the same storage format, whether that\\'s\n",
    "    `torch.strided`{.interpreted-text role=\"attr\"},\n",
    "    `torch.sparse_coo`{.interpreted-text role=\"attr\"}, or\n",
    "    `torch.sparse_csr`{.interpreted-text role=\"attr\"}\n",
    "2.  `data` and `mask` must have the same size, indicated by\n",
    "    `size()`{.interpreted-text role=\"func\"}\n",
    "\n",
    "Sparse COO tensors\n",
    "==================\n",
    "\n",
    "In accordance with Principle \\#1, a sparse COO MaskedTensor is created\n",
    "by passing in two sparse COO tensors, which can be initialized by any of\n",
    "its constructors, for example\n",
    "`torch.sparse_coo_tensor`{.interpreted-text role=\"func\"}.\n",
    "\n",
    "As a recap of [sparse COO\n",
    "tensors](https://pytorch.org/docs/stable/sparse.html#sparse-coo-tensors),\n",
    "the COO format stands for \\\"coordinate format\\\", where the specified\n",
    "elements are stored as tuples of their indices and the corresponding\n",
    "values. That is, the following are provided:\n",
    "\n",
    "-   `indices`: array of size `(ndim, nse)` and dtype `torch.int64`\n",
    "-   `values`: array of size [(nse,)]{.title-ref} with any integer or\n",
    "    floating point dtype\n",
    "\n",
    "where `ndim` is the dimensionality of the tensor and `nse` is the number\n",
    "of specified elements.\n",
    "\n",
    "For both sparse COO and CSR tensors, you can construct a\n",
    "`MaskedTensor`{.interpreted-text role=\"class\"} by doing either:\n",
    "\n",
    "1.  `masked_tensor(sparse_tensor_data, sparse_tensor_mask)`\n",
    "2.  `dense_masked_tensor.to_sparse_coo()` or\n",
    "    `dense_masked_tensor.to_sparse_csr()`\n",
    "\n",
    "The second method is easier to illustrate so we\\'ve shown that below,\n",
    "but for more on the first and the nuances behind the approach, please\n",
    "read the `Sparse COO Appendix <sparse-coo-appendix>`{.interpreted-text\n",
    "role=\"ref\"}.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.masked import masked_tensor\n",
    "import warnings\n",
    "\n",
    "# Disable prototype warnings and such\n",
    "warnings.filterwarnings(action='ignore', category=UserWarning)\n",
    "\n",
    "values = torch.tensor([[0, 0, 3], [4, 0, 5]])\n",
    "mask = torch.tensor([[False, False, True], [False, False, True]])\n",
    "mt = masked_tensor(values, mask)\n",
    "sparse_coo_mt = mt.to_sparse_coo()\n",
    "\n",
    "print(\"mt:\\n\", mt)\n",
    "print(\"mt (sparse coo):\\n\", sparse_coo_mt)\n",
    "print(\"mt data (sparse coo):\\n\", sparse_coo_mt.get_data())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sparse CSR tensors\n",
    "==================\n",
    "\n",
    "Similarly, `MaskedTensor`{.interpreted-text role=\"class\"} also supports\n",
    "the [CSR (Compressed Sparse\n",
    "Row)](https://pytorch.org/docs/stable/sparse.html#sparse-csr-tensor)\n",
    "sparse tensor format. Instead of storing the tuples of the indices like\n",
    "sparse COO tensors, sparse CSR tensors aim to decrease the memory\n",
    "requirements by storing compressed row indices. In particular, a CSR\n",
    "sparse tensor consists of three 1-D tensors:\n",
    "\n",
    "-   `crow_indices`: array of compressed row indices with size\n",
    "    `(size[0] + 1,)`. This array indicates which row a given entry in\n",
    "    values lives in. The last element is the number of specified\n",
    "    elements, while [crow\\_indices\\[i+1\\] -\n",
    "    crow\\_indices\\[i\\]]{.title-ref} indicates the number of specified\n",
    "    elements in row i.\n",
    "-   `col_indices`: array of size `(nnz,)`. Indicates the column indices\n",
    "    for each value.\n",
    "-   `values`: array of size `(nnz,)`. Contains the values of the CSR\n",
    "    tensor.\n",
    "\n",
    "Of note, both sparse COO and CSR tensors are in a\n",
    "[beta](https://pytorch.org/docs/stable/index.html) state.\n",
    "\n",
    "By way of example:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mt_sparse_csr = mt.to_sparse_csr()\n",
    "\n",
    "print(\"mt (sparse csr):\\n\", mt_sparse_csr)\n",
    "print(\"mt data (sparse csr):\\n\", mt_sparse_csr.get_data())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Supported Operations\n",
    "====================\n",
    "\n",
    "Unary\n",
    "-----\n",
    "\n",
    "All [unary\n",
    "operators](https://pytorch.org/docs/master/masked.html#unary-operators)\n",
    "are supported, e.g.:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mt.sin()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Binary\n",
    "======\n",
    "\n",
    "[Binary\n",
    "operators](https://pytorch.org/docs/master/masked.html#unary-operators)\n",
    "are also supported, but the input masks from the two masked tensors must\n",
    "match. For more information on why this decision was made, please find\n",
    "our [MaskedTensor: Advanced Semantics\n",
    "tutorial](https://pytorch.org/tutorials/prototype/maskedtensor_advanced_semantics.html).\n",
    "\n",
    "Please find an example below:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "i = [[0, 1, 1],\n",
    "     [2, 0, 2]]\n",
    "v1 = [3, 4, 5]\n",
    "v2 = [20, 30, 40]\n",
    "m = torch.tensor([True, False, True])\n",
    "\n",
    "s1 = torch.sparse_coo_tensor(i, v1, (2, 3))\n",
    "s2 = torch.sparse_coo_tensor(i, v2, (2, 3))\n",
    "mask = torch.sparse_coo_tensor(i, m, (2, 3))\n",
    "\n",
    "mt1 = masked_tensor(s1, mask)\n",
    "mt2 = masked_tensor(s2, mask)\n",
    "\n",
    "print(\"mt1:\\n\", mt1)\n",
    "print(\"mt2:\\n\", mt2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"torch.div(mt2, mt1):\\n\", torch.div(mt2, mt1))\n",
    "print(\"torch.mul(mt1, mt2):\\n\", torch.mul(mt1, mt2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reductions\n",
    "==========\n",
    "\n",
    "Finally,\n",
    "[reductions](https://pytorch.org/docs/master/masked.html#reductions) are\n",
    "supported:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"mt.sum():\\n\", mt.sum())\n",
    "print(\"mt.sum(dim=1):\\n\", mt.sum(dim=1))\n",
    "print(\"mt.amin():\\n\", mt.amin())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MaskedTensor Helper Methods\n",
    "===========================\n",
    "\n",
    "For convenience, `MaskedTensor`{.interpreted-text role=\"class\"} has a\n",
    "number of methods to help convert between the different layouts and\n",
    "identify the current layout:\n",
    "\n",
    "Setup:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "v = [[3, 0, 0],\n",
    "     [0, 4, 5]]\n",
    "m = [[True, False, False],\n",
    "     [False, True, True]]\n",
    "\n",
    "mt = masked_tensor(torch.tensor(v), torch.tensor(m))\n",
    "mt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MaskedTensor.to_sparse_coo()`{.interpreted-text role=\"meth\"} /\n",
    "`MaskedTensor.to_sparse_csr()`{.interpreted-text role=\"meth\"} /\n",
    "`MaskedTensor.to_dense()`{.interpreted-text role=\"meth\"} to help convert\n",
    "between the different layouts.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mt_sparse_coo = mt.to_sparse_coo()\n",
    "mt_sparse_csr = mt.to_sparse_csr()\n",
    "mt_dense = mt_sparse_coo.to_dense()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MaskedTensor.is_sparse`{.interpreted-text role=\"meth\"} \\-- this will\n",
    "check if the `MaskedTensor`{.interpreted-text role=\"class\"}\\'s layout\n",
    "matches any of the supported sparse layouts (currently COO and CSR).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"mt_dense.is_sparse: \", mt_dense.is_sparse)\n",
    "print(\"mt_sparse_coo.is_sparse: \", mt_sparse_coo.is_sparse)\n",
    "print(\"mt_sparse_csr.is_sparse: \", mt_sparse_csr.is_sparse)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MaskedTensor.is_sparse_coo()`{.interpreted-text role=\"meth\"}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"mt_dense.is_sparse_coo(): \", mt_dense.is_sparse_coo())\n",
    "print(\"mt_sparse_coo.is_sparse_coo: \", mt_sparse_coo.is_sparse_coo())\n",
    "print(\"mt_sparse_csr.is_sparse_coo: \", mt_sparse_csr.is_sparse_coo())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MaskedTensor.is_sparse_csr()`{.interpreted-text role=\"meth\"}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"mt_dense.is_sparse_csr(): \", mt_dense.is_sparse_csr())\n",
    "print(\"mt_sparse_coo.is_sparse_csr: \", mt_sparse_coo.is_sparse_csr())\n",
    "print(\"mt_sparse_csr.is_sparse_csr: \", mt_sparse_csr.is_sparse_csr())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Appendix\n",
    "========\n",
    "\n",
    "Sparse COO Construction {#sparse-coo-appendix}\n",
    "-----------------------\n",
    "\n",
    "Recall in our `original example <sparse-coo-tensors>`{.interpreted-text\n",
    "role=\"ref\"}, we created a `MaskedTensor`{.interpreted-text role=\"class\"}\n",
    "and then converted it to a sparse COO MaskedTensor with\n",
    "`MaskedTensor.to_sparse_coo`{.interpreted-text role=\"meth\"}.\n",
    "\n",
    "Alternatively, we can also construct a sparse COO MaskedTensor directly\n",
    "by passing in two sparse COO tensors:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "values = torch.tensor([[0, 0, 3], [4, 0, 5]]).to_sparse()\n",
    "mask = torch.tensor([[False, False, True], [False, False, True]]).to_sparse()\n",
    "mt = masked_tensor(values, mask)\n",
    "\n",
    "print(\"values:\\n\", values)\n",
    "print(\"mask:\\n\", mask)\n",
    "print(\"mt:\\n\", mt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of using `torch.Tensor.to_sparse`{.interpreted-text\n",
    "role=\"meth\"}, we can also create the sparse COO tensors directly, which\n",
    "brings us to a warning:\n",
    "\n",
    "<div style=\"background-color: #e94f3b; color: #fff; font-weight: 700; padding-left: 10px; padding-top: 5px; padding-bottom: 5px\"><strong>WARNING:</strong></div>\n",
    "\n",
    "<div style=\"background-color: #f3f4f7; padding-left: 10px; padding-top: 10px; padding-bottom: 10px; padding-right: 10px\">\n",
    "\n",
    "<p>When using a function like <code>MaskedTensor.to_sparse_coo</code> (analogous to <code>Tensor.to_sparse</code>),if the user does not specify the indices like in the above example,then the 0 values will be \"unspecified\" by default.</p>\n",
    "\n",
    "</div>\n",
    "\n",
    "Below, we explicitly specify the 0\\'s:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "i = [[0, 1, 1],\n",
    "     [2, 0, 2]]\n",
    "v = [3, 4, 5]\n",
    "m = torch.tensor([True, False, True])\n",
    "values = torch.sparse_coo_tensor(i, v, (2, 3))\n",
    "mask = torch.sparse_coo_tensor(i, m, (2, 3))\n",
    "mt2 = masked_tensor(values, mask)\n",
    "\n",
    "print(\"values:\\n\", values)\n",
    "print(\"mask:\\n\", mask)\n",
    "print(\"mt2:\\n\", mt2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `mt` and `mt2` look identical on the surface, and in the vast\n",
    "majority of operations, will yield the same result. But this brings us\n",
    "to a detail on the implementation:\n",
    "\n",
    "`data` and `mask` \\-- only for sparse MaskedTensors \\-- can have a\n",
    "different number of elements (`nnz`{.interpreted-text role=\"func\"}) **at\n",
    "creation**, but the indices of `mask` must then be a subset of the\n",
    "indices of `data`. In this case, `data` will assume the shape of `mask`\n",
    "by `data = data.sparse_mask(mask)`; in other words, any of the elements\n",
    "in `data` that are not `True` in `mask` (that is, not specified) will be\n",
    "thrown away.\n",
    "\n",
    "Therefore, under the hood, the data looks slightly different; `mt2` has\n",
    "the \\\"4\\\" value masked out and `mt` is completely without it. Their\n",
    "underlying data has different shapes, which would make operations like\n",
    "`mt + mt2` invalid.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"mt data:\\n\", mt.get_data())\n",
    "print(\"mt2 data:\\n\", mt2.get_data())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sparse CSR Construction {#sparse-csr-appendix}\n",
    "=======================\n",
    "\n",
    "We can also construct a sparse CSR MaskedTensor using sparse CSR\n",
    "tensors, and like the example above, this results in a similar treatment\n",
    "under the hood.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "crow_indices = torch.tensor([0, 2, 4])\n",
    "col_indices = torch.tensor([0, 1, 0, 1])\n",
    "values = torch.tensor([1, 2, 3, 4])\n",
    "mask_values = torch.tensor([True, False, False, True])\n",
    "\n",
    "csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.double)\n",
    "mask = torch.sparse_csr_tensor(crow_indices, col_indices, mask_values, dtype=torch.bool)\n",
    "mt = masked_tensor(csr, mask)\n",
    "\n",
    "print(\"mt:\\n\", mt)\n",
    "print(\"mt data:\\n\", mt.get_data())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conclusion\n",
    "==========\n",
    "\n",
    "In this tutorial, we have introduced how to use\n",
    "`MaskedTensor`{.interpreted-text role=\"class\"} with sparse COO and CSR\n",
    "formats and discussed some of the subtleties under the hood in case\n",
    "users decide to access the underlying data structures directly. Sparse\n",
    "storage formats and masked semantics indeed have strong synergies, so\n",
    "much so that they are sometimes used as proxies for each other (as we\n",
    "will see in the next tutorial). In the future, we certainly plan to\n",
    "invest and continue developing in this direction.\n",
    "\n",
    "Further Reading\n",
    "===============\n",
    "\n",
    "To continue learning more, you can find our [Efficiently writing\n",
    "\\\"sparse\\\" semantics for Adagrad with MaskedTensor\n",
    "tutorial](https://pytorch.org/tutorials/prototype/maskedtensor_adagrad.html)\n",
    "to see an example of how MaskedTensor can simplify existing workflows\n",
    "with native masking semantics.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}