notebooks/api/0.8/04-jax-example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SYFT_VERSION = \">=0.8.2.b0,<0.9\"\n",
    "package_string = f'\"syft{SYFT_VERSION}\"'\n",
    "# %pip install {package_string} -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# third party\n",
    "import haiku as hk\n",
    "import jax\n",
    "from jax import random\n",
    "\n",
    "# syft absolute\n",
    "import syft as sy\n",
    "\n",
    "sy.requires(SYFT_VERSION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "node = sy.orchestra.launch(name=\"test-domain-1\", port=\"auto\", dev_mode=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "domain_client = node.login(email=\"info@openmined.org\", password=\"changethis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "key = random.PRNGKey(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_data = random.uniform(key, shape=(4, 28, 28, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert round(train_data.sum()) == 1602"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train = sy.ActionObject.from_obj(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "type(train.syft_action_data), train.id, train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_domain_obj = domain_client.api.services.action.set(train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class MLP(hk.Module):\n",
    "    def __init__(self, out_dims, name=None):\n",
    "        super().__init__(name=name)\n",
    "        self.out_dims = out_dims\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = x.reshape((x.shape[0], -1))\n",
    "        x = hk.Linear(128)(x)\n",
    "        x = jax.nn.relu(x)\n",
    "        x = hk.Linear(self.out_dims)(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "def _forward_fn_linear1(x):\n",
    "    module = MLP(out_dims=10)\n",
    "    return module(x)\n",
    "\n",
    "\n",
    "model = hk.transform(_forward_fn_linear1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "weights = model.init(key, train.syft_action_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert isinstance(weights, dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "w = sy.ActionObject.from_obj(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "type(w.syft_action_data), w.id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "weight_domain_obj = domain_client.api.services.action.set(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@sy.syft_function(\n",
    "    input_policy=sy.ExactMatch(weights=weight_domain_obj.id, data=train_domain_obj.id),\n",
    "    output_policy=sy.SingleExecutionExactOutput(),\n",
    ")\n",
    "def train_mlp(weights, data):\n",
    "    # third party\n",
    "    import haiku as hk\n",
    "    import jax\n",
    "\n",
    "    class MLP(hk.Module):\n",
    "        def __init__(self, out_dims, name=None):\n",
    "            super().__init__(name=name)\n",
    "            self.out_dims = out_dims\n",
    "\n",
    "        def __call__(self, x):\n",
    "            x = x.reshape((x.shape[0], -1))\n",
    "            x = hk.Linear(128)(x)\n",
    "            x = jax.nn.relu(x)\n",
    "            x = hk.Linear(self.out_dims)(x)\n",
    "            return x\n",
    "\n",
    "    def _forward_fn_linear1(x):\n",
    "        module = MLP(out_dims=10)\n",
    "        return module(x)\n",
    "\n",
    "    model = hk.transform(_forward_fn_linear1)\n",
    "    rng_key = jax.random.PRNGKey(42)\n",
    "    output = model.apply(params=weights, x=data, rng=rng_key)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pointer = train_mlp(weights=weight_domain_obj, data=train_domain_obj)\n",
    "output = pointer.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert round(output.sum(), 2) == -0.86"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "request = domain_client.code.request_code_execution(train_mlp)\n",
    "request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "request.approve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "domain_client._api = None\n",
    "_ = domain_client.api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result_ptr = domain_client.code.train_mlp(weights=w.id, data=train.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = result_ptr.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert round(float(result.sum()), 2) == -0.86"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if node.node_type.value == \"python\":\n",
    "    node.land()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}