add SGD robustIV example

Author: apoorvalal

jax_gmm.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -164,7 +164,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -175,7 +175,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -193,6 +193,31 @@
     "    return wt"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Array(6.418394, dtype=float32)"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "sim = ec.data.kang_schafer.Simulation(size=1000)\n",
+    "w, y, X = sim.treatment, sim.outcome, np.c_[np.repeat(1, 1000), sim.covariates]\n",
+    "X0, X1 = X[w == 0, :], X[w == 1, :].mean(0)\n",
+    "b_start = np.random.rand(X.shape[1])\n",
+    "eb_moment(b_start, X0, X1)\n",
+    "# wt = ebwt(X0, X1)\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 4,
@@ -305,7 +330,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,