{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Proof splitting\n",
                "\n",
                "Here we showcase how to split a larger circuit into multiple smaller proofs. This is useful if you want to prove over multiple machines, or if you want to split a proof into multiple parts to reduce the memory requirements.\n",
                "\n",
                "We showcase how to do this in the case where:\n",
                "- intermediate calculations need to be kept secret (but not blinded !) and we need to use the low overhead kzg commitment scheme detailed [here](https://blog.ezkl.xyz/post/commits/) to stitch the circuits together.\n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "First we import the necessary dependencies and set up logging to be as informative as possible. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# check if notebook is in colab\n",
                "try:\n",
                "    # install ezkl\n",
                "    import google.colab\n",
                "    import subprocess\n",
                "    import sys\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"ezkl\"])\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"tf2onnx\"])\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"onnx\"])\n",
                "\n",
                "# rely on local installation of ezkl if the notebook is not in colab\n",
                "except:\n",
                "    pass\n",
                "\n",
                "# make sure you have the dependencies required here already installed\n",
                "import ezkl\n",
                "import os\n",
                "import json\n",
                "import time\n",
                "import random\n",
                "import logging\n",
                "\n",
                "import tensorflow as tf\n",
                "from tensorflow.keras.optimizers.legacy import Adam\n",
                "from tensorflow.keras.layers import *\n",
                "from tensorflow.keras.models import Model\n",
                "from tensorflow.keras.datasets import mnist\n",
                "\n",
                "# uncomment for more descriptive logging \n",
                "# FORMAT = '%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s'\n",
                "# logging.basicConfig(format=FORMAT)\n",
                "# logging.getLogger().setLevel(logging.INFO)\n",
                "\n",
                "# Can we build a simple GAN that can produce all 10 mnist digits?"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
                "x_train, x_test = [x/255.0 for x in [x_train, x_test]]\n",
                "y_train, y_test = [tf.keras.utils.to_categorical(x) for x in [y_train, y_test]]"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "opt = Adam()\n",
                "ZDIM = 100\n",
                "\n",
                "# discriminator\n",
                "# 0 if it's fake, 1 if it's real\n",
                "x = in1 = Input((28,28))\n",
                "x = Reshape((28,28,1))(x)\n",
                "\n",
                "x = Conv2D(64, (5,5), padding='same', strides=(2,2))(x)\n",
                "x = BatchNormalization()(x)\n",
                "x = ELU()(x)\n",
                "\n",
                "x = Conv2D(128, (5,5), padding='same', strides=(2,2))(x)\n",
                "x = BatchNormalization()(x)\n",
                "x = ELU()(x)\n",
                "\n",
                "x = Flatten()(x)\n",
                "x = Dense(128)(x)\n",
                "x = BatchNormalization()(x)\n",
                "x = ELU()(x)\n",
                "x = Dense(1, activation='sigmoid')(x)\n",
                "dm = Model(in1, x)\n",
                "dm.compile(opt, 'binary_crossentropy')\n",
                "dm.summary()\n",
                "\n",
                "# generator, output digits\n",
                "x = in1 = Input((ZDIM,))\n",
                "\n",
                "x = Dense(7*7*64)(x)\n",
                "x = BatchNormalization()(x)\n",
                "x = ELU()(x)\n",
                "x = Reshape((7,7,64))(x)\n",
                "\n",
                "x = Conv2DTranspose(128, (5,5), strides=(2,2), padding='same')(x)\n",
                "x = BatchNormalization()(x)\n",
                "x = ELU()(x)\n",
                "\n",
                "x = Conv2DTranspose(1, (5,5), strides=(2,2), padding='same')(x)\n",
                "x = Activation('sigmoid')(x)\n",
                "x = Reshape((28,28))(x)\n",
                "\n",
                "gm = Model(in1, x)\n",
                "gm.compile('adam', 'mse')\n",
                "gm.summary()\n",
                "\n",
                "# GAN\n",
                "dm.trainable = False\n",
                "x = dm(gm.output)\n",
                "tm = Model(gm.input, x)\n",
                "tm.compile(opt, 'binary_crossentropy')\n",
                "\n",
                "dlosses, glosses = [], []"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import numpy as np\n",
                "from matplotlib.pyplot import figure, imshow, show\n",
                "\n",
                "BS = 256\n",
                "\n",
                "# GAN training loop\n",
                "# make larger if you want it to look better\n",
                "for i in range(1):\n",
                "  # train discriminator\n",
                "  dm.trainable = True\n",
                "  real_i = x_train[np.random.choice(x_train.shape[0], BS)]\n",
                "  fake_i = gm.predict_on_batch(np.random.normal(0,1,size=(BS,ZDIM)))\n",
                "  dloss_r = dm.train_on_batch(real_i, np.ones(BS))\n",
                "  dloss_f = dm.train_on_batch(fake_i, np.zeros(BS))\n",
                "  dloss = (dloss_r + dloss_f)/2\n",
                "\n",
                "  # train generator\n",
                "  dm.trainable = False\n",
                "  gloss_0 = tm.train_on_batch(np.random.normal(0,1,size=(BS,ZDIM)), np.ones(BS))\n",
                "  gloss_1 = tm.train_on_batch(np.random.normal(0,1,size=(BS,ZDIM)), np.ones(BS))\n",
                "  gloss = (gloss_0 + gloss_1)/2\n",
                "\n",
                "  if i%50 == 0:\n",
                "    print(\"%4d: dloss:%8.4f   gloss:%8.4f\" % (i, dloss, gloss))\n",
                "  dlosses.append(dloss)\n",
                "  glosses.append(gloss)\n",
                "    \n",
                "  if i%250 == 0:\n",
                "    \n",
                "    figure(figsize=(16,16))\n",
                "    imshow(np.concatenate(gm.predict(np.random.normal(size=(10,ZDIM))), axis=1))\n",
                "    show()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from matplotlib.pyplot import plot, legend\n",
                "figure(figsize=(8,8))\n",
                "plot(dlosses[100:], label=\"Discriminator Loss\")\n",
                "plot(glosses[100:], label=\"Generator Loss\")\n",
                "legend()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "x = []\n",
                "for i in range(10):\n",
                "  x.append(np.concatenate(gm.predict(np.random.normal(size=(10,ZDIM))), axis=1))\n",
                "imshow(np.concatenate(x, axis=0))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now we export the _generator_ to onnx"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "import numpy as np\n",
                "import tf2onnx\n",
                "import tensorflow as tf\n",
                "import json\n",
                "\n",
                "# split the model into 3 parts\n",
                "gm2 = tf.keras.models.Sequential(gm.layers[0:4])\n",
                "gm3 = tf.keras.models.Sequential(gm.layers[4:11])\n",
                "gm4 = tf.keras.models.Sequential(gm.layers[11:])\n",
                "\n",
                "# After training, export to onnx (network.onnx) and create a data file (input.json)\n",
                "x = 0.1*np.random.rand(1,*[1, ZDIM])\n",
                "inter_x1 = gm2(x[0])\n",
                "inter_x2 = gm3(inter_x1)\n",
                "\n",
                "output_path = \"network_split_0.onnx\"\n",
                "spec = tf.TensorSpec([1, ZDIM], tf.float32, name='input_0')\n",
                "tf2onnx.convert.from_keras(gm2, input_signature=[spec], inputs_as_nchw=['input_0'], opset=12, output_path=output_path)\n",
                "output_path = \"network_split_1.onnx\"\n",
                "spec = tf.TensorSpec(inter_x1.shape, tf.float32, name='elu1')\n",
                "tf2onnx.convert.from_keras(gm3, input_signature=[spec], inputs_as_nchw=['input_1'], opset=12, output_path=output_path)\n",
                "output_path = \"network_split_2.onnx\"\n",
                "spec = tf.TensorSpec(inter_x2.shape, tf.float32, name='elu2')\n",
                "tf2onnx.convert.from_keras(gm4, input_signature=[spec], inputs_as_nchw=['input_2'], opset=12, output_path=output_path)\n",
                "\n",
                "data_array = x.reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_data = [data_array])\n",
                "inter_x1 = inter_x1.numpy().reshape([-1]).tolist()\n",
                "inter_x2 = inter_x2.numpy().reshape([-1]).tolist()\n",
                "data_2 = dict(input_data = [inter_x1])\n",
                "data_3 = dict(input_data = [inter_x2])\n",
                "\n",
                "    # Serialize data into file:\n",
                "data_path = os.path.join('gan_input_0.json')\n",
                "json.dump( data, open(data_path, 'w' ))\n",
                "data_path = os.path.join('gan_input_1.json')\n",
                "json.dump( data_2, open(data_path, 'w' ))\n",
                "data_path = os.path.join('gan_input_2.json')\n",
                "json.dump( data_3, open(data_path, 'w' ))\n",
                "\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "###  KZG commitment intermediate calculations\n",
                "\n",
                "the visibility parameters are:\n",
                "- `input_visibility`: \"kzgcommit\"\n",
                "- `param_visibility`: \"public\"\n",
                "- `output_visibility`: kzgcommit"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import ezkl\n",
                "\n",
                "srs_path = os.path.join('kzg.srs')\n",
                "\n",
                "run_args = ezkl.PyRunArgs()\n",
                "run_args.input_visibility = \"kzgcommit\"\n",
                "run_args.param_visibility = \"fixed\"\n",
                "run_args.output_visibility = \"kzgcommit\"\n",
                "run_args.variables = [(\"batch_size\", 1)]\n",
                "run_args.input_scale = 0\n",
                "run_args.param_scale = 0\n",
                "run_args.logrows = 18\n",
                "\n",
                "ezkl.get_srs(logrows=run_args.logrows)\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Circuit compilation"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 20,
            "metadata": {},
            "outputs": [],
            "source": [
                "# iterate over each submodel gen-settings, compile circuit and setup zkSNARK\n",
                "\n",
                "def setup(i):\n",
                "    print(\"Setting up split model \"+str(i))\n",
                "    # file names\n",
                "    model_path = os.path.join('network_split_'+str(i)+'.onnx')\n",
                "    settings_path = os.path.join('settings_split_'+str(i)+'.json')\n",
                "    data_path =  os.path.join('gan_input_'+str(i)+'.json')\n",
                "    compiled_model_path = os.path.join('network_split_'+str(i)+'.compiled')\n",
                "    pk_path = os.path.join('test_split_'+str(i)+'.pk')\n",
                "    vk_path = os.path.join('test_split_'+str(i)+'.vk')\n",
                "    witness_path = os.path.join('witness_split_'+str(i)+'.json')\n",
                "\n",
                "    if i > 0:\n",
                "         prev_witness_path = os.path.join('witness_split_'+str(i-1)+'.json')\n",
                "         witness = json.load(open(prev_witness_path, 'r'))\n",
                "         data = dict(input_data = witness['outputs'])\n",
                "         # Serialize data into file:\n",
                "         json.dump(data, open(data_path, 'w' ))\n",
                "    else:\n",
                "         data_path = os.path.join('gan_input_0.json')\n",
                "\n",
                "    # generate settings for the current model\n",
                "    res = ezkl.gen_settings(model_path, settings_path, py_run_args=run_args)\n",
                "    res = ezkl.calibrate_settings(data_path, model_path, settings_path, \"resources\", scales=[run_args.input_scale], max_logrows=run_args.logrows)\n",
                "    assert res == True\n",
                "\n",
                "    # load settings and print them to the console\n",
                "    settings = json.load(open(settings_path, 'r'))\n",
                "    settings['run_args']['logrows'] = run_args.logrows\n",
                "    json.dump(settings, open(settings_path, 'w' ))\n",
                "\n",
                "    res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
                "\n",
                "\n",
                "    res = ezkl.setup(\n",
                "         compiled_model_path,\n",
                "         vk_path,\n",
                "         pk_path,\n",
                "         compress_selectors=True,\n",
                "      )\n",
                "\n",
                "    assert res == True\n",
                "    assert os.path.isfile(vk_path)\n",
                "    assert os.path.isfile(pk_path)\n",
                "    res = ezkl.gen_witness(data_path, compiled_model_path, witness_path, vk_path)\n",
                "    run_args.input_scale = settings[\"model_output_scales\"][0]\n",
                "\n",
                "for i in range(3):\n",
                "    setup(i)\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Proof generation"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# GENERATE A PROOF\n",
                "\n",
                "def prove_model(i): \n",
                "    proof_path = os.path.join('proof_split_'+str(i)+'.json')\n",
                "    witness_path = os.path.join('witness_split_'+str(i)+'.json')\n",
                "    compiled_model_path = os.path.join('network_split_'+str(i)+'.compiled')\n",
                "    pk_path = os.path.join('test_split_'+str(i)+'.pk')\n",
                "    vk_path = os.path.join('test_split_'+str(i)+'.vk')\n",
                "    settings_path = os.path.join('settings_split_'+str(i)+'.json')\n",
                "\n",
                "    res = ezkl.prove(\n",
                "            witness_path,\n",
                "            compiled_model_path,\n",
                "            pk_path,\n",
                "            proof_path,\n",
                "            \"for-aggr\",\n",
                "        )\n",
                "\n",
                "    print(res)\n",
                "    assert os.path.isfile(proof_path)\n",
                "\n",
                "    # Verify the proof\n",
                "    if i > 0:\n",
                "        # swap the proof commitments if we are not the first model\n",
                "        prev_witness_path = os.path.join('witness_split_'+str(i-1)+'.json')\n",
                "        prev_witness = json.load(open(prev_witness_path, 'r'))\n",
                "\n",
                "        witness = json.load(open(witness_path, 'r'))\n",
                "\n",
                "        print(prev_witness[\"processed_outputs\"])\n",
                "        print(witness[\"processed_inputs\"])\n",
                "\n",
                "        witness[\"processed_inputs\"] = prev_witness[\"processed_outputs\"]\n",
                "\n",
                "        # now save the witness\n",
                "        with open(witness_path, \"w\") as f:\n",
                "            json.dump(witness, f)\n",
                "\n",
                "        res = ezkl.swap_proof_commitments(proof_path, witness_path)\n",
                "\n",
                "    res = ezkl.verify(\n",
                "            proof_path,\n",
                "            settings_path,\n",
                "            vk_path,\n",
                "        )\n",
                "\n",
                "    assert res == True\n",
                "    print(\"verified\")\n",
                "\n",
                "for i in range(3):\n",
                "    prove_model(i)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "You can also mock aggregate the split proofs into a single proof. This is useful if you want to verify the proof on chain at a lower cost. Here we mock aggregate the proofs to save time. You can use other notebooks to see how to aggregate in full ! "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# now mock aggregate the proofs\n",
                "proofs = []\n",
                "for i in range(3):\n",
                "    proof_path = os.path.join('proof_split_'+str(i)+'.json')\n",
                "    proofs.append(proof_path)\n",
                "\n",
                "ezkl.mock_aggregate(proofs, logrows=22, split_proofs = True)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "ezkl",
            "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.9.15"
        },
        "orig_nbformat": 4
    },
    "nbformat": 4,
    "nbformat_minor": 2
}
