Skip to content

Commit 1ffd986

Browse files
saulshanabrooksaulshanabrook
committed
add mean
1 parent 8880f93 commit 1ffd986

File tree

1 file changed

+35
-5
lines changed

1 file changed

+35
-5
lines changed

‎docs/tutorials/array-api.ipynb

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@
8484
},
8585
{
8686
"cell_type": "code",
87-
"execution_count": 4,
87+
"execution_count": 6,
8888
"metadata": {},
8989
"outputs": [
9090
{
@@ -266,18 +266,30 @@
266266
},
267267
{
268268
"ename": "AttributeError",
269-
"evalue": "module '__main__' has no attribute 'mean'",
269+
"evalue": "Class NDArray does not have method __eq__",
270270
"output_type": "error",
271271
"traceback": [
272272
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
273+
"\u001b[0;31mKeyError\u001b[0m Traceback (most recent call last)",
274+
"File \u001b[0;32m~/p/egg-smol-python/python/egglog/runtime.py:388\u001b[0m, in \u001b[0;36m_special_method\u001b[0;34m(self, __name, *args)\u001b[0m\n\u001b[1;32m 387\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 388\u001b[0m method \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m__egg_decls__\u001b[39m.\u001b[39;49mget_class_decl(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m__egg_typed_expr__\u001b[39m.\u001b[39;49mtp\u001b[39m.\u001b[39;49mname)\u001b[39m.\u001b[39;49mpreserved_methods[__name]\n\u001b[1;32m 389\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m:\n",
275+
"\u001b[0;31mKeyError\u001b[0m: '__eq__'",
276+
"\nDuring handling of the above exception, another exception occurred:\n",
277+
"\u001b[0;31mKeyError\u001b[0m Traceback (most recent call last)",
278+
"File \u001b[0;32m~/p/egg-smol-python/python/egglog/runtime.py:318\u001b[0m, in \u001b[0;36mRuntimeMethod.__post_init__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 317\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 318\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m__egg_fn_decl__ \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m__egg_decls__\u001b[39m.\u001b[39;49mget_function_decl(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m__egg_callable_ref__)\n\u001b[1;32m 319\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m:\n",
279+
"File \u001b[0;32m~/p/egg-smol-python/python/egglog/declarations.py:192\u001b[0m, in \u001b[0;36mModuleDeclarations.get_function_decl\u001b[0;34m(self, ref)\u001b[0m\n\u001b[1;32m 191\u001b[0m \u001b[39mpass\u001b[39;00m\n\u001b[0;32m--> 192\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mKeyError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mFunction \u001b[39m\u001b[39m{\u001b[39;00mref\u001b[39m}\u001b[39;00m\u001b[39m not found\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m 193\u001b[0m \u001b[39melse\u001b[39;00m:\n",
280+
"\u001b[0;31mKeyError\u001b[0m: \"Function MethodRef(class_name='NDArray', method_name='__eq__') not found\"",
281+
"\nDuring handling of the above exception, another exception occurred:\n",
273282
"\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)",
274-
"Cell \u001b[0;32mIn[4], line 682\u001b[0m\n\u001b[1;32m 668\u001b[0m \u001b[39m# Add values for the constants\u001b[39;00m\n\u001b[1;32m 669\u001b[0m egraph\u001b[39m.\u001b[39mregister(\n\u001b[1;32m 670\u001b[0m rewrite(X_arr\u001b[39m.\u001b[39mdtype, runtime_ruleset)\u001b[39m.\u001b[39mto(convert(X\u001b[39m.\u001b[39mdtype, DType)),\n\u001b[1;32m 671\u001b[0m rewrite(y_arr\u001b[39m.\u001b[39mdtype, runtime_ruleset)\u001b[39m.\u001b[39mto(convert(y\u001b[39m.\u001b[39mdtype, DType)),\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 678\u001b[0m rewrite(unique_values(y_arr)\u001b[39m.\u001b[39mshape)\u001b[39m.\u001b[39mto(TupleInt(Int(\u001b[39m3\u001b[39m))),\n\u001b[1;32m 679\u001b[0m )\n\u001b[0;32m--> 682\u001b[0m res \u001b[39m=\u001b[39m fit(X_arr, y_arr)\n\u001b[1;32m 684\u001b[0m \u001b[39m# X_obj, y_obj = egraph.save_object(X), egraph.save_object(y)\u001b[39;00m\n\u001b[1;32m 685\u001b[0m \n\u001b[1;32m 686\u001b[0m \u001b[39m# X_arr = NDArray(X_obj)\u001b[39;00m\n\u001b[1;32m 687\u001b[0m \u001b[39m# y_arr = NDArray(y_obj)\u001b[39;00m\n",
283+
"Cell \u001b[0;32mIn[6], line 700\u001b[0m\n\u001b[1;32m 686\u001b[0m \u001b[39m# Add values for the constants\u001b[39;00m\n\u001b[1;32m 687\u001b[0m egraph\u001b[39m.\u001b[39mregister(\n\u001b[1;32m 688\u001b[0m rewrite(X_arr\u001b[39m.\u001b[39mdtype, runtime_ruleset)\u001b[39m.\u001b[39mto(convert(X\u001b[39m.\u001b[39mdtype, DType)),\n\u001b[1;32m 689\u001b[0m rewrite(y_arr\u001b[39m.\u001b[39mdtype, runtime_ruleset)\u001b[39m.\u001b[39mto(convert(y\u001b[39m.\u001b[39mdtype, DType)),\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 696\u001b[0m rewrite(unique_values(y_arr)\u001b[39m.\u001b[39mshape)\u001b[39m.\u001b[39mto(TupleInt(Int(\u001b[39m3\u001b[39m))),\n\u001b[1;32m 697\u001b[0m )\n\u001b[0;32m--> 700\u001b[0m res \u001b[39m=\u001b[39m fit(X_arr, y_arr)\n\u001b[1;32m 702\u001b[0m \u001b[39m# X_obj, y_obj = egraph.save_object(X), egraph.save_object(y)\u001b[39;00m\n\u001b[1;32m 703\u001b[0m \n\u001b[1;32m 704\u001b[0m \u001b[39m# X_arr = NDArray(X_obj)\u001b[39;00m\n\u001b[1;32m 705\u001b[0m \u001b[39m# y_arr = NDArray(y_obj)\u001b[39;00m\n",
275284
"Cell \u001b[0;32mIn[1], line 15\u001b[0m, in \u001b[0;36mfit\u001b[0;34m(X, y)\u001b[0m\n\u001b[1;32m 13\u001b[0m \u001b[39mwith\u001b[39;00m config_context(array_api_dispatch\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m):\n\u001b[1;32m 14\u001b[0m lda \u001b[39m=\u001b[39m LinearDiscriminantAnalysis(n_components\u001b[39m=\u001b[39m\u001b[39m2\u001b[39m)\n\u001b[0;32m---> 15\u001b[0m X_r2 \u001b[39m=\u001b[39m lda\u001b[39m.\u001b[39;49mfit(X, y)\u001b[39m.\u001b[39mtransform(X)\n\u001b[1;32m 16\u001b[0m \u001b[39mreturn\u001b[39;00m X_r2\n\u001b[1;32m 18\u001b[0m target_names \u001b[39m=\u001b[39m iris\u001b[39m.\u001b[39mtarget_names\n",
276285
"File \u001b[0;32m/usr/local/Caskroom/miniconda/base/envs/egg-smol-python/lib/python3.10/site-packages/sklearn/base.py:1151\u001b[0m, in \u001b[0;36m_fit_context.<locals>.decorator.<locals>.wrapper\u001b[0;34m(estimator, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1144\u001b[0m estimator\u001b[39m.\u001b[39m_validate_params()\n\u001b[1;32m 1146\u001b[0m \u001b[39mwith\u001b[39;00m config_context(\n\u001b[1;32m 1147\u001b[0m skip_parameter_validation\u001b[39m=\u001b[39m(\n\u001b[1;32m 1148\u001b[0m prefer_skip_nested_validation \u001b[39mor\u001b[39;00m global_skip_validation\n\u001b[1;32m 1149\u001b[0m )\n\u001b[1;32m 1150\u001b[0m ):\n\u001b[0;32m-> 1151\u001b[0m \u001b[39mreturn\u001b[39;00m fit_method(estimator, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
277286
"File \u001b[0;32m/usr/local/Caskroom/miniconda/base/envs/egg-smol-python/lib/python3.10/site-packages/sklearn/discriminant_analysis.py:629\u001b[0m, in \u001b[0;36mLinearDiscriminantAnalysis.fit\u001b[0;34m(self, X, y)\u001b[0m\n\u001b[1;32m 623\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcovariance_estimator \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 624\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\n\u001b[1;32m 625\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mcovariance estimator \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 626\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mis not supported \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 627\u001b[0m \u001b[39m\"\u001b[39m\u001b[39mwith svd solver. Try another solver\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m 628\u001b[0m )\n\u001b[0;32m--> 629\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_solve_svd(X, y)\n\u001b[1;32m 630\u001b[0m \u001b[39melif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39msolver \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mlsqr\u001b[39m\u001b[39m\"\u001b[39m:\n\u001b[1;32m 631\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_solve_lstsq(\n\u001b[1;32m 632\u001b[0m X,\n\u001b[1;32m 633\u001b[0m y,\n\u001b[1;32m 634\u001b[0m shrinkage\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mshrinkage,\n\u001b[1;32m 635\u001b[0m covariance_estimator\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcovariance_estimator,\n\u001b[1;32m 636\u001b[0m )\n",
278287
"File \u001b[0;32m/usr/local/Caskroom/miniconda/base/envs/egg-smol-python/lib/python3.10/site-packages/sklearn/discriminant_analysis.py:501\u001b[0m, in \u001b[0;36mLinearDiscriminantAnalysis._solve_svd\u001b[0;34m(self, X, y)\u001b[0m\n\u001b[1;32m 498\u001b[0m n_samples, n_features \u001b[39m=\u001b[39m X\u001b[39m.\u001b[39mshape\n\u001b[1;32m 499\u001b[0m n_classes \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mclasses_\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m]\n\u001b[0;32m--> 501\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmeans_ \u001b[39m=\u001b[39m _class_means(X, y)\n\u001b[1;32m 502\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstore_covariance:\n\u001b[1;32m 503\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mcovariance_ \u001b[39m=\u001b[39m _class_cov(X, y, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mpriors_)\n",
279-
"File \u001b[0;32m/usr/local/Caskroom/miniconda/base/envs/egg-smol-python/lib/python3.10/site-packages/sklearn/discriminant_analysis.py:121\u001b[0m, in \u001b[0;36m_class_means\u001b[0;34m(X, y)\u001b[0m\n\u001b[1;32m 119\u001b[0m \u001b[39mprint\u001b[39m(classes\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m])\n\u001b[1;32m 120\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(classes\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m]):\n\u001b[0;32m--> 121\u001b[0m means[i, :] \u001b[39m=\u001b[39m xp\u001b[39m.\u001b[39;49mmean(X[y \u001b[39m==\u001b[39m i], axis\u001b[39m=\u001b[39m\u001b[39m0\u001b[39m)\n\u001b[1;32m 122\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 123\u001b[0m \u001b[39m# TODO: Explore the choice of using bincount + add.at as it seems sub optimal\u001b[39;00m\n\u001b[1;32m 124\u001b[0m \u001b[39m# from a performance-wise\u001b[39;00m\n\u001b[1;32m 125\u001b[0m cnt \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mbincount(y)\n",
280-
"\u001b[0;31mAttributeError\u001b[0m: module '__main__' has no attribute 'mean'"
288+
"File \u001b[0;32m/usr/local/Caskroom/miniconda/base/envs/egg-smol-python/lib/python3.10/site-packages/sklearn/discriminant_analysis.py:121\u001b[0m, in \u001b[0;36m_class_means\u001b[0;34m(X, y)\u001b[0m\n\u001b[1;32m 119\u001b[0m \u001b[39mprint\u001b[39m(classes\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m])\n\u001b[1;32m 120\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(classes\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m]):\n\u001b[0;32m--> 121\u001b[0m means[i, :] \u001b[39m=\u001b[39m xp\u001b[39m.\u001b[39mmean(X[y \u001b[39m==\u001b[39;49m i], axis\u001b[39m=\u001b[39m\u001b[39m0\u001b[39m)\n\u001b[1;32m 122\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 123\u001b[0m \u001b[39m# TODO: Explore the choice of using bincount + add.at as it seems sub optimal\u001b[39;00m\n\u001b[1;32m 124\u001b[0m \u001b[39m# from a performance-wise\u001b[39;00m\n\u001b[1;32m 125\u001b[0m cnt \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mbincount(y)\n",
289+
"File \u001b[0;32m~/p/egg-smol-python/python/egglog/runtime.py:390\u001b[0m, in \u001b[0;36m_special_method\u001b[0;34m(self, __name, *args)\u001b[0m\n\u001b[1;32m 388\u001b[0m method \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m__egg_decls__\u001b[39m.\u001b[39mget_class_decl(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m__egg_typed_expr__\u001b[39m.\u001b[39mtp\u001b[39m.\u001b[39mname)\u001b[39m.\u001b[39mpreserved_methods[__name]\n\u001b[1;32m 389\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m:\n\u001b[0;32m--> 390\u001b[0m \u001b[39mreturn\u001b[39;00m RuntimeMethod(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m__egg_decls__, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m__egg_typed_expr__, __name)(\u001b[39m*\u001b[39margs)\n\u001b[1;32m 391\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m 392\u001b[0m \u001b[39mreturn\u001b[39;00m method(\u001b[39mself\u001b[39m, \u001b[39m*\u001b[39margs)\n",
290+
"File \u001b[0;32m<string>:6\u001b[0m, in \u001b[0;36m__init__\u001b[0;34m(self, __egg_decls__, __egg_typed_expr__, __egg_method_name__)\u001b[0m\n",
291+
"File \u001b[0;32m~/p/egg-smol-python/python/egglog/runtime.py:320\u001b[0m, in \u001b[0;36mRuntimeMethod.__post_init__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 318\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m__egg_fn_decl__ \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m__egg_decls__\u001b[39m.\u001b[39mget_function_decl(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m__egg_callable_ref__)\n\u001b[1;32m 319\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mKeyError\u001b[39;00m:\n\u001b[0;32m--> 320\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mAttributeError\u001b[39;00m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mClass \u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mclass_name\u001b[39m}\u001b[39;00m\u001b[39m does not have method \u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m__egg_method_name__\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n",
292+
"\u001b[0;31mAttributeError\u001b[0m: Class NDArray does not have method __eq__"
281293
]
282294
}
283295
],
@@ -766,6 +778,21 @@
766778
"converter(Device, OptionalDevice, lambda x: OptionalDevice.some(x))\n",
767779
"\n",
768780
"\n",
781+
"@egraph.class_\n",
782+
"class OptionalTupleInt(Expr):\n",
783+
" none: ClassVar[OptionalTupleInt]\n",
784+
"\n",
785+
" @classmethod\n",
786+
" def some(cls, value: TupleInt) -> OptionalTupleInt:\n",
787+
" ...\n",
788+
"\n",
789+
"\n",
790+
"converter(type(None), OptionalTupleInt, lambda x: OptionalTupleInt.none)\n",
791+
"converter(TupleInt, OptionalTupleInt, lambda x: OptionalTupleInt.some(x))\n",
792+
"converter(int, OptionalTupleInt, lambda x: OptionalTupleInt.some(TupleInt(Int(x))))\n",
793+
"\n",
794+
"\n",
795+
"\n",
769796
"\n",
770797
"\n",
771798
"@egraph.function\n",
@@ -898,6 +925,9 @@
898925
"@egraph.function\n",
899926
"def zeros(shape: TupleInt, dtype: OptionalDType = OptionalDType.none, device: OptionalDevice = OptionalDevice.none) -> NDArray:\n",
900927
" ...\n",
928+
"@egraph.function\n",
929+
"def mean(x: NDArray, axis: OptionalTupleInt = OptionalTupleInt.none) -> NDArray: ...\n",
930+
"\n",
901931
"\n",
902932
"linalg = sys.modules[__name__]\n",
903933
"\n",

0 commit comments

Comments
 (0)