benchmarks/benchmarks.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3a7f5936",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jaxtyping import Float, Array\n",
    "from typing import Any\n",
    "import jax.tree_util as jtu\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "07d25a9f",
   "metadata": {},
   "source": [
    "# Mytree code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "536f66e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/danieldodd/miniconda3/lib/python3.10/site-packages/flax/core/frozen_dict.py:169: FutureWarning: jax.tree_util.register_keypaths is deprecated, and will be removed in a future release. Please use `register_pytree_with_keys()` instead.\n",
      "  jax.tree_util.register_keypaths(\n"
     ]
    }
   ],
   "source": [
    "from mytree import Mytree, param_field, Softplus\n",
    "\n",
    "class Mytree_SubFoo(Mytree):\n",
    "    a: Float[Array, \"...\"] = param_field(bijector=Softplus)\n",
    "    b: Float[Array, \"...\"] = param_field(bijector=Softplus)\n",
    "\n",
    "    def __init__(self, a, b):\n",
    "        self.a = a\n",
    "        self.b = b\n",
    "\n",
    "class Mytree_Foo(Mytree):\n",
    "    b: list[Mytree_SubFoo]\n",
    "    a: Float[Array, \"...\"] = param_field(bijector=Softplus)\n",
    "\n",
    "    def __init__(self, a, b):\n",
    "        self.a = a\n",
    "        self.b = b"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e298d9b9",
   "metadata": {},
   "source": [
    "# Pytree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1ce3a220",
   "metadata": {},
   "outputs": [],
   "source": [
    "from simple_pytree import Pytree, static_field\n",
    "\n",
    "class Param(Pytree):\n",
    "    value: Any\n",
    "    bijector: Any = static_field()\n",
    "    trainable: Any = static_field()\n",
    "\n",
    "    def __init__(self, value, bijector=Softplus, trainable=True):\n",
    "        self.value = value\n",
    "        self.bijector = bijector\n",
    "        self.trainable = trainable\n",
    "\n",
    "def _is_param(x):\n",
    "    return isinstance(x, Param)\n",
    "\n",
    "def _resolve_bijector_forward(x):\n",
    "    if _is_param(x):\n",
    "\n",
    "        return x.replace(value = x.bijector.forward(x.value))\n",
    "    else:\n",
    "        return x\n",
    "    \n",
    "def _resolve_bijector_inverse(x):\n",
    "    if _is_param(x):\n",
    "        return x.replace(value = x.bijector.inverse(x.value))\n",
    "    else:\n",
    "        return x\n",
    "\n",
    "class Module(Pytree):    \n",
    "    def constrain(self):\n",
    "        return jtu.tree_map(_resolve_bijector_forward, self, is_leaf=_is_param)\n",
    "    \n",
    "    def unconstrain(self):\n",
    "        return jtu.tree_map(_resolve_bijector_inverse, self, is_leaf=_is_param)\n",
    "\n",
    " \n",
    "\n",
    "class Pytree_SubFoo(Module):\n",
    "    a: Param \n",
    "    b: Param \n",
    "\n",
    "    def __init__(self, a: Float[Array, \"...\"] , b: Float[Array, \"...\"]):\n",
    "        self.a = Param(a)\n",
    "        self.b = Param(b)\n",
    "\n",
    "\n",
    "class Pytree_Foo(Module):\n",
    "    b: list[Pytree_SubFoo]\n",
    "    a: Param\n",
    "\n",
    "    def __init__(self, b: list[Pytree_SubFoo], a:  Float[Array, \"...\"]):\n",
    "        self.b = b\n",
    "        self.a = Param(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dc5a35a",
   "metadata": {},
   "source": [
    "# Performance comparison:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0f5978c5",
   "metadata": {},
   "source": [
    "Run on a M1 Pro CPU.\n",
    "\n",
    "- **Initialisation**: is faster for mytree.\n",
    "- **Transformations**: is faster for mytree.\n",
    "- **Replacing attributes**: is faster for mytree implimentation.\n",
    "\n",
    "Completeing further benchmarks is on the todo list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8db39ca8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " mytree:\n",
      "52.1 ms ± 1.18 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "1.02 s ± 35 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "1.52 µs ± 15.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n",
      "\n",
      " pytree:\n",
      "58 ms ± 2.65 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "1.08 s ± 20.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "2.05 µs ± 76.2 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "pytree_classes = [Pytree_Foo, Pytree_SubFoo]\n",
    "mytree_classes = [Mytree_Foo, Mytree_SubFoo]\n",
    "\n",
    "\n",
    "def init_tree(tree, subtree):\n",
    "    return tree([subtree(jnp.array([1.0] * 10000), jnp.array([1.0] * 10000))]*10000, jnp.array([3.0] * 10000))\n",
    "\n",
    "for name, implimentation in zip([\"mytree\", \"pytree\"], [mytree_classes, pytree_classes]):\n",
    "\n",
    "    print(f\"\\n {name}:\")\n",
    "    %timeit init_tree(*implimentation)\n",
    "    foo = init_tree(*implimentation)\n",
    "    %timeit foo.constrain()\n",
    "    %timeit foo.replace(a=123)"
   ]
  }
 ],
 "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.10.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "c99f523e7e05658ebf273d5a51641edbaa579bc207415a2fcfcb80446d5d6a90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}