Metadata via trace

benchmarks/benchmarks.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "id": "3a7f5936",
    "metadata": {},
    "outputs": [],
@@ -24,10 +24,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "id": "536f66e1",
    "metadata": {},
-   "outputs": [],
+   "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",
@@ -59,7 +68,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "id": "1ce3a220",
    "metadata": {},
    "outputs": [],
@@ -135,7 +144,7 @@
    "source": [
     "Run on a M1 Pro CPU.\n",
     "\n",
-    "- **Initialisation**: is faster for mytree, despite it unpacking metadata, and working out what attributes are leaves of the nested pytree structure.\n",
+    "- **Initialisation**: is faster for mytree.\n",
     "- **Transformations**: is faster for mytree.\n",
     "- **Replacing attributes**: is faster for mytree implimentation.\n",
     "\n",
@@ -144,7 +153,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 4,
    "id": "8db39ca8",
    "metadata": {},
    "outputs": [
@@ -154,14 +163,14 @@
      "text": [
       "\n",
       " mytree:\n",
-      "50.7 ms ± 1.05 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
-      "862 ms ± 13.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
-      "1.34 µs ± 31.3 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\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",
-      "51.8 ms ± 626 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
-      "893 ms ± 24.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
-      "1.7 µs ± 44.4 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\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"
      ]
     }
    ],

mytree/bijectors.py

@@ -2,19 +2,17 @@
 
 __all__ = ["Bijector", "Identity", "Softplus"]
 
-import importlib
-import jax.numpy as jnp
-from typing import Callable
 from dataclasses import dataclass
+from typing import Callable
+
+import jax.numpy as jnp
 from simple_pytree import Pytree, static_field
 
+
 @dataclass
 class Bijector(Pytree):
-    forward: Callable = static_field()
-    inverse: Callable = static_field()
-
-def __init__(self, forward: Callable, inverse: Callable) -> None:
-    """Initialise the bijector.
+    """
+    Create a bijector.
 
     Args:
         forward(Callable): The forward transformation.
@@ -23,13 +21,14 @@ def __init__(self, forward: Callable, inverse: Callable) -> None:
     Returns:
         Bijector: A bijector.
     """
-    self.forward = forward
-    self.inverse = inverse
+
+    forward: Callable = static_field()
+    inverse: Callable = static_field()
 
 
 Identity = Bijector(forward=lambda x: x, inverse=lambda x: x)
 
 Softplus = Bijector(
-    forward=lambda x: jnp.log(1 + jnp.exp(x)),
+    forward=lambda x: jnp.log(1.0 + jnp.exp(x)),
     inverse=lambda x: jnp.log(jnp.exp(x) - 1.0),
 )