docs/tutorials/shape_constraints.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7765UFHoyGx6"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KsOkK8O69PyT"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKQpW0JqQQmY"
      },
      "source": [
        "# Shape Constraints with Tensorflow Lattice\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r61fkA2i9Y3_"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/lattice/tutorials/shape_constraints\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/lattice/blob/master/docs/tutorials/shape_constraints.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/lattice/blob/master/docs/tutorials/shape_constraints.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/lattice/docs/tutorials/shape_constraints.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2plcL3iTVjsp"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial is an overview of the constraints and regularizers provided by the TensorFlow Lattice (TFL) library. Here we use TFL premade models on synthetic datasets, but note that everything in this tutorial can also be done with models constructed from TFL Keras layers.\n",
        "\n",
        "Before proceeding, make sure your runtime has all required packages installed (as imported in the code cells below)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x769lI12IZXB"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fbBVAR6UeRN5"
      },
      "source": [
        "Installing TF Lattice package:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bpXjJKpSd3j4"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install -U tensorflow tf-keras tensorflow-lattice pydot graphviz\n",
        "!pip install -U tensorflow_decision_forests"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jSVl9SHTeSGX"
      },
      "source": [
        "Importing required packages:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "iY6awAl058TV"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_lattice as tfl\n",
        "import tensorflow_decision_forests as tfdf\n",
        "\n",
        "from IPython.core.pylabtools import figsize\n",
        "import functools\n",
        "import logging\n",
        "import matplotlib\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import sys\n",
        "import tempfile\n",
        "logging.disable(sys.maxsize)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8dsfk2oNlakY"
      },
      "outputs": [],
      "source": [
        "# Use Keras 2.\n",
        "version_fn = getattr(tf.keras, \"version\", None)\n",
        "if version_fn and version_fn().startswith(\"3.\"):\n",
        "  import tf_keras as keras\n",
        "else:\n",
        "  keras = tf.keras"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7TmBk_IGgJF0"
      },
      "source": [
        "Default values used in this guide:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kQHPyPsPUF92"
      },
      "outputs": [],
      "source": [
        "NUM_EPOCHS = 1000\n",
        "BATCH_SIZE = 64\n",
        "LEARNING_RATE=0.01"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjR7D8Ag3z0d"
      },
      "source": [
        "## Training Dataset for Ranking Restaurants"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a1YetzbdFOij"
      },
      "source": [
        "Imagine a simplified scenario where we want to determine whether or not users will click on a restaurant search result. The task is to predict the clickthrough rate (CTR) given input features:\n",
        "- Average rating (`avg_rating`): a numeric feature with values in the range [1,5].\n",
        "- Number of reviews (`num_reviews`): a numeric feature with values capped at 200, which we use as a measure of trendiness.\n",
        "- Dollar rating (`dollar_rating`): a categorical feature with string values in the set {\"D\", \"DD\", \"DDD\", \"DDDD\"}.\n",
        "\n",
        "Here we create a synthetic dataset where the true CTR is given by the formula:\n",
        "$$\n",
        "CTR = 1 / (1 + exp\\{\\mbox{b(dollar_rating)}-\\mbox{avg_rating}\\times log(\\mbox{num_reviews}) /4 \\})\n",
        "$$\n",
        "where $b(\\cdot)$ translates each `dollar_rating` to a baseline value:\n",
        "$$\n",
        "\\mbox{D}\\to 3,\\ \\mbox{DD}\\to 2,\\ \\mbox{DDD}\\to 4,\\ \\mbox{DDDD}\\to 4.5.\n",
        "$$\n",
        "\n",
        "This formula reflects typical user patterns. e.g. given everything else fixed, users prefer restaurants with higher star ratings, and \"\\\\$\\\\$\" restaurants will receive more clicks than \"\\\\$\", followed by \"\\\\$\\\\$\\\\$\" and \"\\\\$\\\\$\\\\$\\\\$\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mKovnyv1jATw"
      },
      "outputs": [],
      "source": [
        "dollar_ratings_vocab = [\"D\", \"DD\", \"DDD\", \"DDDD\"]\n",
        "def click_through_rate(avg_ratings, num_reviews, dollar_ratings):\n",
        "  dollar_rating_baseline = {\"D\": 3, \"DD\": 2, \"DDD\": 4, \"DDDD\": 4.5}\n",
        "  return 1 / (1 + np.exp(\n",
        "      np.array([dollar_rating_baseline[d] for d in dollar_ratings]) -\n",
        "      avg_ratings * np.log1p(num_reviews) / 4))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BPlgRdt6jAbP"
      },
      "source": [
        "Let's take a look at the contour plots of this CTR function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KC5qX_XKmc7g"
      },
      "outputs": [],
      "source": [
        "def color_bar():\n",
        "  bar = matplotlib.cm.ScalarMappable(\n",
        "      norm=matplotlib.colors.Normalize(0, 1, True),\n",
        "      cmap=\"viridis\",\n",
        "  )\n",
        "  bar.set_array([0, 1])\n",
        "  return bar\n",
        "\n",
        "\n",
        "def plot_fns(fns, res=25):\n",
        "  \"\"\"Generates contour plots for a list of (name, fn) functions.\"\"\"\n",
        "  num_reviews, avg_ratings = np.meshgrid(\n",
        "      np.linspace(0, 200, num=res),\n",
        "      np.linspace(1, 5, num=res),\n",
        "  )\n",
        "  figsize(13, 3.5 * len(fns))\n",
        "  fig, axes = plt.subplots(\n",
        "      len(fns), len(dollar_ratings_vocab), sharey=True, layout=\"constrained\"\n",
        "  )\n",
        "  axes = axes.flatten()\n",
        "  axes_index = 0\n",
        "  for fn_name, fn in fns:\n",
        "    for dollar_rating_split in dollar_ratings_vocab:\n",
        "      dollar_ratings = np.repeat(dollar_rating_split, res**2)\n",
        "      values = fn(avg_ratings.flatten(), num_reviews.flatten(), dollar_ratings)\n",
        "      title = \"{}: dollar_rating={}\".format(fn_name, dollar_rating_split)\n",
        "      subplot = axes[axes_index]\n",
        "      axes_index += 1\n",
        "      subplot.contourf(\n",
        "          avg_ratings,\n",
        "          num_reviews,\n",
        "          np.reshape(values, (res, res)),\n",
        "          vmin=0,\n",
        "          vmax=1,\n",
        "      )\n",
        "      subplot.title.set_text(title)\n",
        "      subplot.set(xlabel=\"Average Rating\")\n",
        "      subplot.set(ylabel=\"Number of Reviews\")\n",
        "      subplot.set(xlim=(1, 5))\n",
        "\n",
        "  if len(fns) \u003c= 2:\n",
        "    cax = fig.add_axes([\n",
        "        axes[-1].get_position().x1 + 0.11,\n",
        "        axes[-1].get_position().y0,\n",
        "        0.02,\n",
        "        0.8,\n",
        "    ])\n",
        "    _ = fig.colorbar(color_bar(), cax=cax)\n",
        "\n",
        "\n",
        "plot_fns([(\"CTR\", click_through_rate)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ol91olp3muNN"
      },
      "source": [
        "### Preparing Data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H8BOshZS9xwn"
      },
      "source": [
        "We now need to create our synthetic datasets. We start by generating a simulated dataset of restaurants and their features."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MhqcOPdTT_wj"
      },
      "outputs": [],
      "source": [
        "def sample_restaurants(n):\n",
        "  avg_ratings = np.random.uniform(1.0, 5.0, n)\n",
        "  num_reviews = np.round(np.exp(np.random.uniform(0.0, np.log(200), n)))\n",
        "  dollar_ratings = np.random.choice(dollar_ratings_vocab, n)\n",
        "  ctr_labels = click_through_rate(avg_ratings, num_reviews, dollar_ratings)\n",
        "  return avg_ratings, num_reviews, dollar_ratings, ctr_labels\n",
        "\n",
        "\n",
        "np.random.seed(42)\n",
        "avg_ratings, num_reviews, dollar_ratings, ctr_labels = sample_restaurants(2000)\n",
        "\n",
        "figsize(5, 5)\n",
        "fig, axs = plt.subplots(1, 1, sharey=False, layout=\"constrained\")\n",
        "\n",
        "for rating, marker in [(\"D\", \"o\"), (\"DD\", \"^\"), (\"DDD\", \"+\"), (\"DDDD\", \"x\")]:\n",
        "  plt.scatter(\n",
        "      x=avg_ratings[np.where(dollar_ratings == rating)],\n",
        "      y=num_reviews[np.where(dollar_ratings == rating)],\n",
        "      c=ctr_labels[np.where(dollar_ratings == rating)],\n",
        "      vmin=0,\n",
        "      vmax=1,\n",
        "      marker=marker,\n",
        "      label=rating)\n",
        "plt.xlabel(\"Average Rating\")\n",
        "plt.ylabel(\"Number of Reviews\")\n",
        "plt.legend()\n",
        "plt.xlim((1, 5))\n",
        "plt.title(\"Distribution of restaurants\")\n",
        "_ = fig.colorbar(color_bar(), cax=fig.add_axes([1.05, 0.1, 0.05, 0.85]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tRetsfLv_JSR"
      },
      "source": [
        "Let's produce the training, validation and testing datasets. When a restaurant is viewed in the search results, we can record user's engagement (click or no click) as a sample point.\n",
        "\n",
        "In practice, users often do not go through all search results. This means that users will likely only see restaurants already considered \"good\" by the current ranking model in use. As a result, \"good\" restaurants are more frequently impressed and over-represented in the training datasets. When using more features, the training dataset can have large gaps in \"bad\" parts of the feature space.\n",
        "\n",
        "When the model is used for ranking, it is often evaluated on all relevant results with a more uniform distribution that is not well-represented by the training dataset. A flexible and complicated model might fail in this case due to overfitting the over-represented data points and thus lack generalizability. We handle this issue by applying domain knowledge to add *shape constraints* that guide the model to make reasonable predictions when it cannot pick them up from the training dataset.\n",
        "\n",
        "In this example, the training dataset mostly consists of user interactions with good and popular restaurants. The testing dataset has a uniform distribution to simulate the evaluation setting discussed above. Note that such testing dataset will not be available in a real problem setting."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jS6WOtXQ8jwX"
      },
      "outputs": [],
      "source": [
        "def sample_dataset(n, testing_set):\n",
        "  (avg_ratings, num_reviews, dollar_ratings, ctr_labels) = sample_restaurants(n)\n",
        "  if testing_set:\n",
        "    # Testing has a more uniform distribution over all restaurants.\n",
        "    num_views = np.random.poisson(lam=3, size=n)\n",
        "  else:\n",
        "    # Training/validation datasets have more views on popular restaurants.\n",
        "    num_views = np.random.poisson(lam=ctr_labels * num_reviews / 50.0, size=n)\n",
        "\n",
        "  return pd.DataFrame({\n",
        "      \"avg_rating\": np.repeat(avg_ratings, num_views),\n",
        "      \"num_reviews\": np.repeat(num_reviews, num_views),\n",
        "      \"dollar_rating\": np.repeat(dollar_ratings, num_views),\n",
        "      \"clicked\": np.random.binomial(n=1, p=np.repeat(ctr_labels, num_views)),\n",
        "  })\n",
        "\n",
        "\n",
        "# Generate datasets.\n",
        "np.random.seed(42)\n",
        "data_train = sample_dataset(500, testing_set=False)\n",
        "data_val = sample_dataset(500, testing_set=False)\n",
        "data_test = sample_dataset(500, testing_set=True)\n",
        "\n",
        "ds_train = tfdf.keras.pd_dataframe_to_tf_dataset(\n",
        "    data_train, label=\"clicked\", batch_size=BATCH_SIZE\n",
        ")\n",
        "ds_val = tfdf.keras.pd_dataframe_to_tf_dataset(\n",
        "    data_val, label=\"clicked\", batch_size=BATCH_SIZE\n",
        ")\n",
        "ds_test = tfdf.keras.pd_dataframe_to_tf_dataset(\n",
        "    data_test, label=\"clicked\", batch_size=BATCH_SIZE\n",
        ")\n",
        "\n",
        "# feature_analysis_data is used to find quantiles of featurse.\n",
        "feature_analysis_data = data_train.copy()\n",
        "feature_analysis_data[\"dollar_rating\"] = feature_analysis_data[\n",
        "    \"dollar_rating\"\n",
        "].map({v: i for i, v in enumerate(dollar_ratings_vocab)})\n",
        "feature_analysis_data = dict(feature_analysis_data)\n",
        "\n",
        "# Plotting dataset densities.\n",
        "figsize(12, 5)\n",
        "fig, axs = plt.subplots(1, 2, sharey=False, tight_layout=False)\n",
        "for ax, data, title in [\n",
        "    (axs[0], data_train, \"training\"),\n",
        "    (axs[1], data_test, \"testing\"),\n",
        "]:\n",
        "  _, _, _, density = ax.hist2d(\n",
        "      x=data[\"avg_rating\"],\n",
        "      y=data[\"num_reviews\"],\n",
        "      bins=(np.linspace(1, 5, num=21), np.linspace(0, 200, num=21)),\n",
        "      cmap=\"Blues\",\n",
        "  )\n",
        "  ax.set(xlim=(1, 5))\n",
        "  ax.set(ylim=(0, 200))\n",
        "  ax.set(xlabel=\"Average Rating\")\n",
        "  ax.set(ylabel=\"Number of Reviews\")\n",
        "  ax.title.set_text(\"Density of {} examples\".format(title))\n",
        "  _ = fig.colorbar(density, ax=ax)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qoTrw3FZqvPK"
      },
      "source": [
        "## Fitting Gradient Boosted Trees"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZklNowexE3wB"
      },
      "source": [
        "We first create a few auxillary functions for plotting and calculating validation and test metrics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3BqGqScQzlYf"
      },
      "outputs": [],
      "source": [
        "def pred_fn(model, from_logits, avg_ratings, num_reviews, dollar_rating):\n",
        "  preds = model.predict(\n",
        "      tf.data.Dataset.from_tensor_slices({\n",
        "          \"avg_rating\": avg_ratings,\n",
        "          \"num_reviews\": num_reviews,\n",
        "          \"dollar_rating\": dollar_rating,\n",
        "      }).batch(1),\n",
        "      verbose=0,\n",
        "  )\n",
        "  if from_logits:\n",
        "    preds = tf.math.sigmoid(preds)\n",
        "  return preds\n",
        "\n",
        "\n",
        "def analyze_model(models, from_logits=False, print_metrics=True):\n",
        "  pred_fns = []\n",
        "  for model, name in models:\n",
        "    if print_metrics:\n",
        "      metric = model.evaluate(ds_val, return_dict=True, verbose=0)\n",
        "      print(\"Validation AUC: {}\".format(metric[\"auc\"]))\n",
        "      metric = model.evaluate(ds_test, return_dict=True, verbose=0)\n",
        "      print(\"Testing AUC: {}\".format(metric[\"auc\"]))\n",
        "\n",
        "    pred_fns.append(\n",
        "        (\"{} pCTR\".format(name), functools.partial(pred_fn, model, from_logits))\n",
        "    )\n",
        "\n",
        "  pred_fns.append((\"CTR\", click_through_rate))\n",
        "  plot_fns(pred_fns)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JVef4f8yUUbs"
      },
      "source": [
        "We can fit TensorFlow gradient boosted decision trees on the dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DnPYlRAo2mnQ"
      },
      "outputs": [],
      "source": [
        "gbt_model = tfdf.keras.GradientBoostedTreesModel(\n",
        "    features=[\n",
        "        tfdf.keras.FeatureUsage(name=\"num_reviews\"),\n",
        "        tfdf.keras.FeatureUsage(name=\"avg_rating\"),\n",
        "        tfdf.keras.FeatureUsage(name=\"dollar_rating\"),\n",
        "    ],\n",
        "    exclude_non_specified_features=True,\n",
        "    num_threads=1,\n",
        "    num_trees=32,\n",
        "    max_depth=6,\n",
        "    min_examples=10,\n",
        "    growing_strategy=\"BEST_FIRST_GLOBAL\",\n",
        "    random_seed=42,\n",
        "    temp_directory=tempfile.mkdtemp(),\n",
        ")\n",
        "gbt_model.compile(metrics=[keras.metrics.AUC(name=\"auc\")])\n",
        "gbt_model.fit(ds_train, validation_data=ds_val, verbose=0)\n",
        "analyze_model([(gbt_model, \"GBT\")])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYZtd6YvsNdn"
      },
      "source": [
        "Even though the model has captured the general shape of the true CTR and has decent validation metrics, it has counter-intuitive behavior in several parts of the input space: the estimated CTR decreases as the average rating or number of reviews increase. This is due to a lack of sample points in areas not well-covered by the training dataset. The model simply has no way to deduce the correct behaviour solely from the data.\n",
        "\n",
        "To solve this issue, we enforce the shape constraint that the model must output values monotonically increasing with respect to both the average rating and the number of reviews. We will later see how to implement this in TFL.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uf7WqGooFiEp"
      },
      "source": [
        "## Fitting a DNN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_s2aT3x0E_tF"
      },
      "source": [
        "We can repeat the same steps with a DNN classifier. We can observe a similar pattern: not having enough sample points with small number of reviews results in nonsensical extrapolation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WKZzCY-UkZX-"
      },
      "outputs": [],
      "source": [
        "keras.utils.set_random_seed(42)\n",
        "inputs = {\n",
        "    \"num_reviews\": keras.Input(shape=(1,), dtype=tf.float32),\n",
        "    \"avg_rating\": keras.Input(shape=(1), dtype=tf.float32),\n",
        "    \"dollar_rating\": keras.Input(shape=(1), dtype=tf.string),\n",
        "}\n",
        "inputs_flat = keras.layers.Concatenate()([\n",
        "    inputs[\"num_reviews\"],\n",
        "    inputs[\"avg_rating\"],\n",
        "    keras.layers.StringLookup(\n",
        "        vocabulary=dollar_ratings_vocab,\n",
        "        num_oov_indices=0,\n",
        "        output_mode=\"one_hot\",\n",
        "    )(inputs[\"dollar_rating\"]),\n",
        "])\n",
        "dense_layers = keras.Sequential(\n",
        "    [\n",
        "        keras.layers.Dense(16, activation=\"relu\"),\n",
        "        keras.layers.Dense(16, activation=\"relu\"),\n",
        "        keras.layers.Dense(1, activation=None),\n",
        "    ],\n",
        "    name=\"dense_layers\",\n",
        ")\n",
        "dnn_model = keras.Model(inputs=inputs, outputs=dense_layers(inputs_flat))\n",
        "keras.utils.plot_model(\n",
        "    dnn_model, expand_nested=True, show_layer_names=False, rankdir=\"LR\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6zFqu6wf1I30"
      },
      "outputs": [],
      "source": [
        "dnn_model.compile(\n",
        "    loss=keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "    metrics=[keras.metrics.AUC(from_logits=True, name=\"auc\")],\n",
        "    optimizer=keras.optimizers.Adam(LEARNING_RATE),\n",
        ")\n",
        "dnn_model.fit(ds_train, epochs=200, verbose=0)\n",
        "analyze_model([(dnn_model, \"DNN\")], from_logits=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Avkw-okw7JL"
      },
      "source": [
        "## Shape Constraints"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3ExyethCFBrP"
      },
      "source": [
        "TensorFlow Lattice (TFL) is focused on enforcing shape constraints to safeguard model behavior beyond the training data. These shape constraints are applied to TFL Keras layers. Their details can be found in [our JMLR paper](http://jmlr.org/papers/volume17/15-243/15-243.pdf).\n",
        "\n",
        "In this tutorial we use TF premade models to cover various shape constraints, but note that all these steps can be done with models created from TFL Keras layers.\n",
        "\n",
        "Using TFL premade models also requires:\n",
        "- a *model config*: defining the model architecture and per-feature shape constraints and regularizers.\n",
        "- a *feature analysis dataset*: a dataset used for TFL initialization (feature quantile calcuation).\n",
        "\n",
        "For a more thorough description, please refer to the premade models or the API docs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "anyCM4sCpOSo"
      },
      "source": [
        "### Monotonicity\n",
        "We first address the monotonicity concerns by adding monotonicity shape constraints to the continuous features. We use a calibrated lattice model with added output calibration: each feature is calibrated using categorical or piecewise-linear calibrators, then fed into a lattice model, followed by an output piecewise-linear calibrator.\n",
        "\n",
        "To instruct TFL to enforce shape constraints, we specify the constraints in the *feature configs*. The following code shows how we can require the output to be monotonically increasing with respect to both `num_reviews` and `avg_rating` by setting `monotonicity=\"increasing\"`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hFlkZs5RgFcP"
      },
      "outputs": [],
      "source": [
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"dollar_rating\",\n",
        "            lattice_size=3,\n",
        "            pwl_calibration_num_keypoints=4,\n",
        "            vocabulary_list=dollar_ratings_vocab,\n",
        "            num_buckets=len(dollar_ratings_vocab),\n",
        "        ),\n",
        "    ],\n",
        "    output_calibration=True,\n",
        "    output_initialization=np.linspace(-2, 2, num=5),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GOlzuyQsGre5"
      },
      "source": [
        "We now use the `feature_analysis_data` to find and set the quantile values for the input features. These values can be pre-calculated and set explicitly in the feature config depending on the training pipeline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f-bTmfBnghuX"
      },
      "outputs": [],
      "source": [
        "feature_analysis_data = data_train.copy()\n",
        "feature_analysis_data[\"dollar_rating\"] = feature_analysis_data[\n",
        "    \"dollar_rating\"\n",
        "].map({v: i for i, v in enumerate(dollar_ratings_vocab)})\n",
        "feature_analysis_data = dict(feature_analysis_data)\n",
        "\n",
        "feature_keypoints = tfl.premade_lib.compute_feature_keypoints(\n",
        "    feature_configs=model_config.feature_configs, features=feature_analysis_data\n",
        ")\n",
        "tfl.premade_lib.set_feature_keypoints(\n",
        "    feature_configs=model_config.feature_configs,\n",
        "    feature_keypoints=feature_keypoints,\n",
        "    add_missing_feature_configs=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FCm1lOjmwur_"
      },
      "outputs": [],
      "source": [
        "keras.utils.set_random_seed(42)\n",
        "inputs = {\n",
        "    \"num_reviews\": keras.Input(shape=(1,), dtype=tf.float32),\n",
        "    \"avg_rating\": keras.Input(shape=(1), dtype=tf.float32),\n",
        "    \"dollar_rating\": keras.Input(shape=(1), dtype=tf.string),\n",
        "}\n",
        "ordered_inputs = [\n",
        "    inputs[\"num_reviews\"],\n",
        "    inputs[\"avg_rating\"],\n",
        "    keras.layers.StringLookup(\n",
        "        vocabulary=dollar_ratings_vocab,\n",
        "        num_oov_indices=0,\n",
        "        output_mode=\"int\",\n",
        "    )(inputs[\"dollar_rating\"]),\n",
        "]\n",
        "outputs = tfl.premade.CalibratedLattice(\n",
        "    model_config=model_config, name=\"CalibratedLattice\"\n",
        ")(ordered_inputs)\n",
        "tfl_model_0 = keras.Model(inputs=inputs, outputs=outputs)\n",
        "\n",
        "keras.utils.plot_model(\n",
        "    tfl_model_0, expand_nested=True, show_layer_names=False, rankdir=\"LR\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ubNRBCWW5wQ9"
      },
      "source": [
        "Using a `CalibratedLatticeConfig` creates a premade classifier that first applies a *calibrator* to each input (a piece-wise linear function for numeric features) followed by a *lattice* layer to non-linearly fuse the calibrated features. We have also enabled output piece-wise linear calibration.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Am1OwtzzU7no"
      },
      "outputs": [],
      "source": [
        "tfl_model_0.compile(\n",
        "    loss=keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "    metrics=[keras.metrics.AUC(from_logits=True, name=\"auc\")],\n",
        "    optimizer=keras.optimizers.Adam(LEARNING_RATE),\n",
        ")\n",
        "tfl_model_0.fit(ds_train, epochs=100, verbose=0)\n",
        "analyze_model([(tfl_model_0, \"TFL0\")], from_logits=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7vZ5fShXs504"
      },
      "source": [
        "With the constraints added, the estimated CTR will always increase as the average rating increases or the number of reviews increases. This is done by making sure that the calibrators and the lattice are monotonic."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pSUd6aFlpYz4"
      },
      "source": [
        "### Partial Monotonicity for Categorical Calibration\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CnPiqf4rq6kJ"
      },
      "source": [
        "To use constraints on the third feature, `dollar_rating`, we should recall that categorical features require a slightly different treatment in TFL. Here we enforce the partial monotonicity constraint that outputs for \"DD\" restaurants should be larger than \"D\" restaurants when all other inputs are fixed. This is done using the `monotonicity` setting in the feature config. We also need to use `tfl.premade_lib.set_categorical_monotonicities` to convert the constrains specified in string values into the numerical format understood by the library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FH2ItfsTsE3S"
      },
      "outputs": [],
      "source": [
        "keras.utils.set_random_seed(42)\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_convexity=\"concave\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"dollar_rating\",\n",
        "            lattice_size=3,\n",
        "            pwl_calibration_num_keypoints=4,\n",
        "            vocabulary_list=dollar_ratings_vocab,\n",
        "            num_buckets=len(dollar_ratings_vocab),\n",
        "            monotonicity=[(\"D\", \"DD\")],\n",
        "        ),\n",
        "    ],\n",
        "    output_calibration=True,\n",
        "    output_initialization=np.linspace(-2, 2, num=5),\n",
        ")\n",
        "\n",
        "tfl.premade_lib.set_feature_keypoints(\n",
        "    feature_configs=model_config.feature_configs,\n",
        "    feature_keypoints=feature_keypoints,\n",
        "    add_missing_feature_configs=False,\n",
        ")\n",
        "tfl.premade_lib.set_categorical_monotonicities(model_config.feature_configs)\n",
        "\n",
        "outputs = tfl.premade.CalibratedLattice(\n",
        "    model_config=model_config, name=\"CalibratedLattice\"\n",
        ")(ordered_inputs)\n",
        "tfl_model_1 = keras.Model(inputs=inputs, outputs=outputs)\n",
        "tfl_model_1.compile(\n",
        "    loss=keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "    metrics=[keras.metrics.AUC(from_logits=True, name=\"auc\")],\n",
        "    optimizer=keras.optimizers.Adam(LEARNING_RATE),\n",
        ")\n",
        "tfl_model_1.fit(ds_train, epochs=100, verbose=0)\n",
        "analyze_model([(tfl_model_1, \"TFL1\")], from_logits=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gdIzhYL79_Pp"
      },
      "source": [
        "Here we also plot the predicted CTR of this model conditioned on `dollar_rating`. Notice that all the constraints we required are fulfilled in each of the slices."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J6CP2Ovapiu3"
      },
      "source": [
        "### 2D Shape Constraint: Trust\n",
        "A 5-star rating for a restaurant with only one or two reviews is likely an unreliable rating (the restaurant might not actually be good), whereas a 4-star rating for a restaurant with hundreds of reviews is much more reliable (the restaurant is likely good in this case). We can see that the number of reviews of a restaurant affects how much trust we place in its average rating.\n",
        "\n",
        "We can exercise TFL trust constraints to inform the model that the larger (or smaller) value of one feature indicates more reliance or trust of another feature. This is done by setting `reflects_trust_in` configuration in the feature config."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OA14j0erm6TJ"
      },
      "outputs": [],
      "source": [
        "keras.utils.set_random_seed(42)\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "            # Larger num_reviews indicating more trust in avg_rating.\n",
        "            reflects_trust_in=[\n",
        "                tfl.configs.TrustConfig(\n",
        "                    feature_name=\"avg_rating\", trust_type=\"edgeworth\"\n",
        "                ),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"dollar_rating\",\n",
        "            lattice_size=3,\n",
        "            pwl_calibration_num_keypoints=4,\n",
        "            vocabulary_list=dollar_ratings_vocab,\n",
        "            num_buckets=len(dollar_ratings_vocab),\n",
        "            monotonicity=[(\"D\", \"DD\")],\n",
        "        ),\n",
        "    ],\n",
        "    output_calibration=True,\n",
        "    output_initialization=np.linspace(-2, 2, num=5),\n",
        ")\n",
        "\n",
        "tfl.premade_lib.set_feature_keypoints(\n",
        "    feature_configs=model_config.feature_configs,\n",
        "    feature_keypoints=feature_keypoints,\n",
        "    add_missing_feature_configs=False,\n",
        ")\n",
        "tfl.premade_lib.set_categorical_monotonicities(model_config.feature_configs)\n",
        "\n",
        "outputs = tfl.premade.CalibratedLattice(\n",
        "    model_config=model_config, name=\"CalibratedLattice\"\n",
        ")(ordered_inputs)\n",
        "tfl_model_2 = keras.Model(inputs=inputs, outputs=outputs)\n",
        "tfl_model_2.compile(\n",
        "    loss=keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "    metrics=[keras.metrics.AUC(from_logits=True, name=\"auc\")],\n",
        "    optimizer=keras.optimizers.Adam(LEARNING_RATE),\n",
        ")\n",
        "tfl_model_2.fit(ds_train, epochs=100, verbose=0)\n",
        "analyze_model([(tfl_model_2, \"TFL2\")], from_logits=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "puvP9X8XxyRV"
      },
      "source": [
        "The following plot presents the trained lattice function. Due to the trust constraint, we expect that larger values of calibrated `num_reviews` would force higher slope with respect to calibrated `avg_rating`, resulting in a more significant move in the lattice output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "RounEQebxxnA"
      },
      "outputs": [],
      "source": [
        "lattice_params = tfl_model_2.layers[-1].layers[-2].weights[0].numpy()\n",
        "lat_mesh_x, lat_mesh_y = np.meshgrid(\n",
        "    np.linspace(0, 1, num=3),\n",
        "    np.linspace(0, 1, num=3),\n",
        ")\n",
        "lat_mesh_z = np.reshape(np.asarray(lattice_params[0::3]), (3, 3))\n",
        "\n",
        "figure = plt.figure(figsize=(6, 6))\n",
        "axes = figure.add_subplot(projection=\"3d\")\n",
        "axes.plot_wireframe(lat_mesh_x, lat_mesh_y, lat_mesh_z, color=\"dodgerblue\")\n",
        "plt.legend([\"Lattice Lookup\"])\n",
        "plt.title(\"Trust\")\n",
        "plt.xlabel(\"Calibrated avg_rating\")\n",
        "plt.ylabel(\"Calibrated num_reviews\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RfniRZCHIvfK"
      },
      "source": [
        "### Diminishing Returns\n",
        "[Diminishing returns](https://en.wikipedia.org/wiki/Diminishing_returns) means that the marginal gain of increasing a certain feature value will decrease as we increase the value. In our case we expect that the `num_reviews` feature follows this pattern, so we can configure its calibrator accordingly. Notice that we can decompose diminishing returns into two sufficient conditions:\n",
        "\n",
        "- the calibrator is monotonicially increasing, and\n",
        "- the calibrator is concave (setting `pwl_calibration_convexity=\"concave\"`).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XQrM9BskY-wx"
      },
      "outputs": [],
      "source": [
        "keras.utils.set_random_seed(42)\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_convexity=\"concave\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "            reflects_trust_in=[\n",
        "                tfl.configs.TrustConfig(\n",
        "                    feature_name=\"avg_rating\", trust_type=\"edgeworth\"\n",
        "                ),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"dollar_rating\",\n",
        "            lattice_size=3,\n",
        "            pwl_calibration_num_keypoints=4,\n",
        "            vocabulary_list=dollar_ratings_vocab,\n",
        "            num_buckets=len(dollar_ratings_vocab),\n",
        "            monotonicity=[(\"D\", \"DD\")],\n",
        "        ),\n",
        "    ],\n",
        "    output_calibration=True,\n",
        "    output_initialization=np.linspace(-2, 2, num=5),\n",
        ")\n",
        "\n",
        "tfl.premade_lib.set_feature_keypoints(\n",
        "    feature_configs=model_config.feature_configs,\n",
        "    feature_keypoints=feature_keypoints,\n",
        "    add_missing_feature_configs=False,\n",
        ")\n",
        "tfl.premade_lib.set_categorical_monotonicities(model_config.feature_configs)\n",
        "\n",
        "outputs = tfl.premade.CalibratedLattice(\n",
        "    model_config=model_config, name=\"CalibratedLattice\"\n",
        ")(ordered_inputs)\n",
        "tfl_model_3 = keras.Model(inputs=inputs, outputs=outputs)\n",
        "tfl_model_3.compile(\n",
        "    loss=keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "    metrics=[keras.metrics.AUC(from_logits=True, name=\"auc\")],\n",
        "    optimizer=keras.optimizers.Adam(LEARNING_RATE),\n",
        ")\n",
        "tfl_model_3.fit(\n",
        "    ds_train,\n",
        "    epochs=100,\n",
        "    verbose=0\n",
        ")\n",
        "analyze_model([(tfl_model_3, \"TFL3\")], from_logits=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LSmzHkPUo9u5"
      },
      "source": [
        "Notice how the testing metric improves by adding the concavity constraint. The prediction plot also better resembles the ground truth."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SKe3UHX6pUjw"
      },
      "source": [
        "### Smoothing Calibrators\n",
        "We notice in the prediction curves above that even though the output is monotonic in specified features, the changes in the slopes are abrupt and hard to interpret. That suggests we might want to consider smoothing this calibrator using a regularizer setup in the `regularizer_configs`.\n",
        "\n",
        "Here we apply a `hessian` regularizer to make the calibration more linear. You can also use the `laplacian` regularizer to flatten the calibrator and the `wrinkle` regularizer to reduce changes in the curvature.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CxcCNxhkqC7u"
      },
      "outputs": [],
      "source": [
        "keras.utils.set_random_seed(42)\n",
        "model_config = tfl.configs.CalibratedLatticeConfig(\n",
        "    feature_configs=[\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"num_reviews\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_convexity=\"concave\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "            regularizer_configs=[\n",
        "                tfl.configs.RegularizerConfig(name=\"calib_hessian\", l2=0.5),\n",
        "            ],\n",
        "            reflects_trust_in=[\n",
        "                tfl.configs.TrustConfig(\n",
        "                    feature_name=\"avg_rating\", trust_type=\"edgeworth\"\n",
        "                ),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"avg_rating\",\n",
        "            lattice_size=3,\n",
        "            monotonicity=\"increasing\",\n",
        "            pwl_calibration_num_keypoints=32,\n",
        "            regularizer_configs=[\n",
        "                tfl.configs.RegularizerConfig(name=\"calib_hessian\", l2=0.5),\n",
        "            ],\n",
        "        ),\n",
        "        tfl.configs.FeatureConfig(\n",
        "            name=\"dollar_rating\",\n",
        "            lattice_size=3,\n",
        "            pwl_calibration_num_keypoints=4,\n",
        "            vocabulary_list=dollar_ratings_vocab,\n",
        "            num_buckets=len(dollar_ratings_vocab),\n",
        "            monotonicity=[(\"D\", \"DD\")],\n",
        "        ),\n",
        "    ],\n",
        "    output_calibration=True,\n",
        "    output_initialization=np.linspace(-2, 2, num=5),\n",
        "    regularizer_configs=[\n",
        "        tfl.configs.RegularizerConfig(name=\"calib_hessian\", l2=0.1),\n",
        "    ],\n",
        ")\n",
        "\n",
        "tfl.premade_lib.set_feature_keypoints(\n",
        "    feature_configs=model_config.feature_configs,\n",
        "    feature_keypoints=feature_keypoints,\n",
        "    add_missing_feature_configs=False,\n",
        ")\n",
        "tfl.premade_lib.set_categorical_monotonicities(model_config.feature_configs)\n",
        "\n",
        "outputs = tfl.premade.CalibratedLattice(\n",
        "    model_config=model_config, name=\"CalibratedLattice\"\n",
        ")(ordered_inputs)\n",
        "tfl_model_4 = keras.Model(inputs=inputs, outputs=outputs)\n",
        "tfl_model_4.compile(\n",
        "    loss=keras.losses.BinaryCrossentropy(from_logits=True),\n",
        "    metrics=[keras.metrics.AUC(from_logits=True, name=\"auc\")],\n",
        "    optimizer=keras.optimizers.Adam(LEARNING_RATE),\n",
        ")\n",
        "tfl_model_4.fit(ds_train, epochs=100, verbose=0)\n",
        "analyze_model([(tfl_model_4, \"TFL4\")], from_logits=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HHpp4goLvuPi"
      },
      "source": [
        "The calibrators are now smooth, and the overall estimated CTR better matches the ground truth. This is reflected both in the testing metric and in the contour plots."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TLOGDrYY0hH7"
      },
      "source": [
        "Here you can see the results of each step as we added domain-specific constraints and regularizers to the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nUEuihX815ix"
      },
      "outputs": [],
      "source": [
        "analyze_model(\n",
        "    [\n",
        "        (tfl_model_0, \"TFL0\"),\n",
        "        (tfl_model_1, \"TFL1\"),\n",
        "        (tfl_model_2, \"TFL2\"),\n",
        "        (tfl_model_3, \"TFL3\"),\n",
        "        (tfl_model_4, \"TFL4\"),\n",
        "    ],\n",
        "    from_logits=True,\n",
        "    print_metrics=False,\n",
        ")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "shape_constraints.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}