Add new equinox example (#4)

* new equinox

* format

Author: lockwo

examples/02_combining_metaclasses.ipynb

@@ -2,40 +2,58 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "value: 1.0\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Array(2., dtype=float32, weak_type=True)"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "# import equinox as eqx\n",
-    "# import ihoop\n",
-    "# import jax\n",
+    "import equinox as eqx\n",
+    "import ihoop\n",
+    "import jax\n",
+    "\n",
     "\n",
-    "# # works with https://github.com/patrick-kidger/equinox/pull/1028\n",
-    "# class _StrictEqxMeta(ihoop.strict._StrictMeta, eqx._module._module._ModuleMeta):\n",
-    "#     def __new__(mcs, name, bases, namespace, **kwargs):\n",
-    "#         return super().__new__(mcs, name, bases, namespace, **kwargs)\n",
+    "class _StrictEqxMeta(ihoop.strict._StrictMeta, eqx._module._module._ModuleMeta):\n",
+    "    def __new__(mcs, name, bases, namespace, **kwargs):\n",
+    "        return super().__new__(mcs, name, bases, namespace, **kwargs)\n",
     "\n",
     "\n",
-    "# class StrictModule(eqx.Module, ihoop.Strict, metaclass=_StrictEqxMeta):\n",
-    "#     pass\n",
+    "class StrictModule(eqx.Module, ihoop.Strict, metaclass=_StrictEqxMeta):\n",
+    "    pass\n",
     "\n",
     "\n",
-    "# class Test(StrictModule):\n",
-    "#     x: float\n",
+    "class Test(StrictModule):\n",
+    "    x: float\n",
     "\n",
-    "#     def __init__(self):\n",
-    "#         self.x = 1.0\n",
+    "    def __init__(self):\n",
+    "        self.x = 1.0\n",
     "\n",
     "\n",
-    "# t = Test()\n",
-    "# print(\"value:\", t.x)\n",
+    "t = Test()\n",
+    "print(\"value:\", t.x)\n",
     "\n",
     "\n",
-    "# def f(x):\n",
-    "#     return x.x**2\n",
+    "def f(x):\n",
+    "    return x.x**2\n",
     "\n",
     "\n",
-    "# jax.grad(f)(t).x"
+    "jax.jit(jax.grad(f))(t).x"
    ]
   },
   {

ihoop/strict.py

@@ -37,10 +37,10 @@ def __new__(
         """
         Runs when a class inheriting from Strict is defined.
         """
-        # just check the initial letter as keyword, that way we can have multiple
+        # just check the initial letters as keyword, that way we can have multiple
         # Strict classes that could resolve metaclass conflicts
         is_defining_strict_itself = (
-            name == "Strict" and namespace.get("__module__") == mcs.__module__
+            name[:6] == "Strict" and namespace.get("__module__") == mcs.__module__
         )
 
         if not is_defining_strict_itself:

pyproject.toml

@@ -24,7 +24,7 @@ build-backend = "setuptools.build_meta"
 
 [project.optional-dependencies]
 examples = [
-    "equinox>=0.12.0"
+    "equinox>=0.13.0"
 ]
 testing = [
     "pytest==7.2.0",