_NDArray_1 = NDArray.var("X")\n",
"assume_dtype(_NDArray_1, DType.float64)\n",
"assume_shape(_NDArray_1, TupleInt(Int(1000000)) + TupleInt(Int(20)))\n",
"assume_isfinite(_NDArray_1)\n",
"_NDArray_2 = NDArray.var("y")\n",
"assume_dtype(_NDArray_2, DType.int64)\n",
"assume_shape(_NDArray_2, TupleInt(Int(1000000)))\n",
"assume_value_one_of(_NDArray_2, TupleValue(Value.int(Int(0))) + TupleValue(Value.int(Int(1))))\n",
"_NDArray_3 = astype(\n",
" NDArray.vector(TupleValue(sum(_NDArray_2 == NDArray.scalar(Value.int(Int(0)))).to_value()) + TupleValue(sum(_NDArray_2 == NDArray.scalar(Value.int(Int(1)))).to_value())),\n",
" DType.float64,\n",
") / NDArray.scalar(Value.float(Float(1000000.0)))\n",
"_NDArray_4 = zeros(TupleInt(Int(2)) + TupleInt(Int(20)), OptionalDType.some(DType.float64), OptionalDevice.some(_NDArray_1.device))\n",
"_MultiAxisIndexKey_1 = MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice()))\n",
"_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(0))) + _MultiAxisIndexKey_1)\n",
"_NDArray_5 = _NDArray_1[ndarray_index(_NDArray_2 == NDArray.scalar(Value.int(Int(0))))]\n",
"_OptionalIntOrTuple_1 = OptionalIntOrTuple.some(IntOrTuple.int(Int(0)))\n",
"_NDArray_4[_IndexKey_1] = sum(_NDArray_5, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_5.shape[Int(0)]))\n",
"_IndexKey_2 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(1))) + _MultiAxisIndexKey_1)\n",
"_NDArray_6 = _NDArray_1[ndarray_index(_NDArray_2 == NDArray.scalar(Value.int(Int(1))))]\n",
"_NDArray_4[_IndexKey_2] = sum(_NDArray_6, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_6.shape[Int(0)]))\n",
"_NDArray_7 = concat(TupleNDArray(_NDArray_5 - _NDArray_4[_IndexKey_1]) + TupleNDArray(_NDArray_6 - _NDArray_4[_IndexKey_2]), OptionalInt.some(Int(0)))\n",
"_NDArray_8 = square(_NDArray_7 - expand_dims(sum(_NDArray_7, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_7.shape[Int(0)]))))\n",
"_NDArray_9 = sqrt(sum(_NDArray_8, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_8.shape[Int(0)])))\n",
"_NDArray_10 = copy(_NDArray_9)\n",
"_NDArray_10[ndarray_index(_NDArray_9 == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(Value.float(Float(1.0)))\n",
"_TupleNDArray_1 = svd(sqrt(NDArray.scalar(Value.float(Float.rational(Rational(1, 999998))))) * (_NDArray_7 / _NDArray_10), FALSE)\n",
"_Slice_1 = Slice(OptionalInt.none, OptionalInt.some(sum(astype(_TupleNDArray_1[Int(1)] > NDArray.scalar(Value.float(Float(0.0001))), DType.int32)).to_value().to_int))\n",
"_NDArray_11 = (_TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(_Slice_1)) + _MultiAxisIndexKey_1)] / _NDArray_10).T / _TupleNDArray_1[\n",
" Int(1)\n",
"][IndexKey.slice(_Slice_1)]\n",
"_TupleNDArray_2 = svd(\n",
" (sqrt((NDArray.scalar(Value.int(Int(1000000))) * _NDArray_3) * NDArray.scalar(Value.float(Float(1.0)))) * (_NDArray_4 - (_NDArray_3 @ _NDArray_4)).T).T @ _NDArray_11, FALSE\n",
")\n",
"(\n",
" (_NDArray_1 - (_NDArray_3 @ _NDArray_4))\n",
" @ (\n",
" _NDArray_11\n",
" @ _TupleNDArray_2[Int(2)].T[\n",
" IndexKey.multi_axis(\n",
" _MultiAxisIndexKey_1\n",
" + MultiAxisIndexKey(\n",
" MultiAxisIndexKeyItem.slice(\n",
" Slice(\n",
" OptionalInt.none,\n",
" OptionalInt.some(\n",
" sum(astype(_TupleNDArray_2[Int(1)] > (NDArray.scalar(Value.float(Float(0.0001))) * _TupleNDArray_2[Int(1)][IndexKey.int(Int(0))]), DType.int32))\n",
" .to_value()\n",
" .to_int\n",
" ),\n",
" )\n",
" )\n",
" )\n",
" )\n",
" ]\n",
" )\n",
")[IndexKey.multi_axis(_MultiAxisIndexKey_1 + MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(1))))))]\n",
" \n"
],
"text/latex": [
"\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}1} \\PY{o}{=} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{var}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{X}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
"\\PY{n}{assume\\PYZus{}dtype}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{,} \\PY{n}{DType}\\PY{o}{.}\\PY{n}{float64}\\PY{p}{)}\n",
"\\PY{n}{assume\\PYZus{}shape}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1000000}\\PY{p}{)}\\PY{p}{)} \\PY{o}{+} \\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{20}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{assume\\PYZus{}isfinite}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}2} \\PY{o}{=} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{var}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
"\\PY{n}{assume\\PYZus{}dtype}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{,} \\PY{n}{DType}\\PY{o}{.}\\PY{n}{int64}\\PY{p}{)}\n",
"\\PY{n}{assume\\PYZus{}shape}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1000000}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{assume\\PYZus{}value\\PYZus{}one\\PYZus{}of}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{,} \\PY{n}{TupleValue}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{+} \\PY{n}{TupleValue}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}3} \\PY{o}{=} \\PY{n}{astype}\\PY{p}{(}\n",
" \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{vector}\\PY{p}{(}\\PY{n}{TupleValue}\\PY{p}{(}\\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2} \\PY{o}{==} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{o}{.}\\PY{n}{to\\PYZus{}value}\\PY{p}{(}\\PY{p}{)}\\PY{p}{)} \\PY{o}{+} \\PY{n}{TupleValue}\\PY{p}{(}\\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2} \\PY{o}{==} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{o}{.}\\PY{n}{to\\PYZus{}value}\\PY{p}{(}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,}\n",
" \\PY{n}{DType}\\PY{o}{.}\\PY{n}{float64}\\PY{p}{,}\n",
"\\PY{p}{)} \\PY{o}{/} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{p}{(}\\PY{l+m+mf}{1000000.0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{=} \\PY{n}{zeros}\\PY{p}{(}\\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{)} \\PY{o}{+} \\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{20}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,} \\PY{n}{OptionalDType}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{DType}\\PY{o}{.}\\PY{n}{float64}\\PY{p}{)}\\PY{p}{,} \\PY{n}{OptionalDevice}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{o}{.}\\PY{n}{device}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}MultiAxisIndexKey\\PYZus{}1} \\PY{o}{=} \\PY{n}{MultiAxisIndexKey}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{Slice}\\PY{p}{(}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}IndexKey\\PYZus{}1} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{+} \\PY{n}{\\PYZus{}MultiAxisIndexKey\\PYZus{}1}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}5} \\PY{o}{=} \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2} \\PY{o}{==} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\n",
"\\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1} \\PY{o}{=} \\PY{n}{OptionalIntOrTuple}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{IntOrTuple}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}1}\\PY{p}{]} \\PY{o}{=} \\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}5}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)} \\PY{o}{/} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}5}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{+} \\PY{n}{\\PYZus{}MultiAxisIndexKey\\PYZus{}1}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}6} \\PY{o}{=} \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2} \\PY{o}{==} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2}\\PY{p}{]} \\PY{o}{=} \\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}6}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)} \\PY{o}{/} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}6}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{=} \\PY{n}{concat}\\PY{p}{(}\\PY{n}{TupleNDArray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}5} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}1}\\PY{p}{]}\\PY{p}{)} \\PY{o}{+} \\PY{n}{TupleNDArray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}6} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2}\\PY{p}{]}\\PY{p}{)}\\PY{p}{,} \\PY{n}{OptionalInt}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}8} \\PY{o}{=} \\PY{n}{square}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{\\PYZhy{}} \\PY{n}{expand\\PYZus{}dims}\\PY{p}{(}\\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)} \\PY{o}{/} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}9} \\PY{o}{=} \\PY{n}{sqrt}\\PY{p}{(}\\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}8}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)} \\PY{o}{/} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}8}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}10} \\PY{o}{=} \\PY{n}{copy}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}9}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}10}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}9} \\PY{o}{==} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]} \\PY{o}{=} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{p}{(}\\PY{l+m+mf}{1.0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1} \\PY{o}{=} \\PY{n}{svd}\\PY{p}{(}\\PY{n}{sqrt}\\PY{p}{(}\\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{o}{.}\\PY{n}{rational}\\PY{p}{(}\\PY{n}{Rational}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{,} \\PY{l+m+mi}{999998}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{*} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}10}\\PY{p}{)}\\PY{p}{,} \\PY{n}{FALSE}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}Slice\\PYZus{}1} \\PY{o}{=} \\PY{n}{Slice}\\PY{p}{(}\\PY{n}{OptionalInt}\\PY{o}{.}\\PY{n}{none}\\PY{p}{,} \\PY{n}{OptionalInt}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{astype}\\PY{p}{(}\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{]} \\PY{o}{\\PYZgt{}} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{p}{(}\\PY{l+m+mf}{0.0001}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,} \\PY{n}{DType}\\PY{o}{.}\\PY{n}{int32}\\PY{p}{)}\\PY{p}{)}\\PY{o}{.}\\PY{n}{to\\PYZus{}value}\\PY{p}{(}\\PY{p}{)}\\PY{o}{.}\\PY{n}{to\\PYZus{}int}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}NDArray\\PYZus{}11} \\PY{o}{=} \\PY{p}{(}\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{]}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{\\PYZus{}Slice\\PYZus{}1}\\PY{p}{)}\\PY{p}{)} \\PY{o}{+} \\PY{n}{\\PYZus{}MultiAxisIndexKey\\PYZus{}1}\\PY{p}{)}\\PY{p}{]} \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}10}\\PY{p}{)}\\PY{o}{.}\\PY{n}{T} \\PY{o}{/} \\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1}\\PY{p}{[}\n",
" \\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\n",
"\\PY{p}{]}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{\\PYZus{}Slice\\PYZus{}1}\\PY{p}{)}\\PY{p}{]}\n",
"\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}2} \\PY{o}{=} \\PY{n}{svd}\\PY{p}{(}\n",
" \\PY{p}{(}\\PY{n}{sqrt}\\PY{p}{(}\\PY{p}{(}\\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1000000}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{*} \\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{)} \\PY{o}{*} \\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{p}{(}\\PY{l+m+mf}{1.0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{*} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{\\PYZhy{}} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3} \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{)}\\PY{p}{)}\\PY{o}{.}\\PY{n}{T}\\PY{p}{)}\\PY{o}{.}\\PY{n}{T} \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}11}\\PY{p}{,} \\PY{n}{FALSE}\n",
"\\PY{p}{)}\n",
"\\PY{p}{(}\n",
" \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1} \\PY{o}{\\PYZhy{}} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3} \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{)}\\PY{p}{)}\n",
" \\PY{o}{@} \\PY{p}{(}\n",
" \\PY{n}{\\PYZus{}NDArray\\PYZus{}11}\n",
" \\PY{o}{@} \\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}2}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{]}\\PY{o}{.}\\PY{n}{T}\\PY{p}{[}\n",
" \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\n",
" \\PY{n}{\\PYZus{}MultiAxisIndexKey\\PYZus{}1}\n",
" \\PY{o}{+} \\PY{n}{MultiAxisIndexKey}\\PY{p}{(}\n",
" \\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\n",
" \\PY{n}{Slice}\\PY{p}{(}\n",
" \\PY{n}{OptionalInt}\\PY{o}{.}\\PY{n}{none}\\PY{p}{,}\n",
" \\PY{n}{OptionalInt}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\n",
" \\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{astype}\\PY{p}{(}\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}2}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{]} \\PY{o}{\\PYZgt{}} \\PY{p}{(}\\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{p}{(}\\PY{l+m+mf}{0.0001}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{*} \\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}2}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{]}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)}\\PY{p}{,} \\PY{n}{DType}\\PY{o}{.}\\PY{n}{int32}\\PY{p}{)}\\PY{p}{)}\n",
" \\PY{o}{.}\\PY{n}{to\\PYZus{}value}\\PY{p}{(}\\PY{p}{)}\n",
" \\PY{o}{.}\\PY{n}{to\\PYZus{}int}\n",
" \\PY{p}{)}\\PY{p}{,}\n",
" \\PY{p}{)}\n",
" \\PY{p}{)}\n",
" \\PY{p}{)}\n",
" \\PY{p}{)}\n",
" \\PY{p}{]}\n",
" \\PY{p}{)}\n",
"\\PY{p}{)}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{\\PYZus{}MultiAxisIndexKey\\PYZus{}1} \\PY{o}{+} \\PY{n}{MultiAxisIndexKey}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{Slice}\\PY{p}{(}\\PY{n}{OptionalInt}\\PY{o}{.}\\PY{n}{none}\\PY{p}{,} \\PY{n}{OptionalInt}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\n",
"\\end{Verbatim}\n"
],
"text/plain": [
"_NDArray_1 = NDArray.var(\"X\")\n",
"assume_dtype(_NDArray_1, DType.float64)\n",
"assume_shape(_NDArray_1, TupleInt(Int(1000000)) + TupleInt(Int(20)))\n",
"assume_isfinite(_NDArray_1)\n",
"_NDArray_2 = NDArray.var(\"y\")\n",
"assume_dtype(_NDArray_2, DType.int64)\n",
"assume_shape(_NDArray_2, TupleInt(Int(1000000)))\n",
"assume_value_one_of(_NDArray_2, TupleValue(Value.int(Int(0))) + TupleValue(Value.int(Int(1))))\n",
"_NDArray_3 = astype(\n",
" NDArray.vector(TupleValue(sum(_NDArray_2 == NDArray.scalar(Value.int(Int(0)))).to_value()) + TupleValue(sum(_NDArray_2 == NDArray.scalar(Value.int(Int(1)))).to_value())),\n",
" DType.float64,\n",
") / NDArray.scalar(Value.float(Float(1000000.0)))\n",
"_NDArray_4 = zeros(TupleInt(Int(2)) + TupleInt(Int(20)), OptionalDType.some(DType.float64), OptionalDevice.some(_NDArray_1.device))\n",
"_MultiAxisIndexKey_1 = MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice()))\n",
"_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(0))) + _MultiAxisIndexKey_1)\n",
"_NDArray_5 = _NDArray_1[ndarray_index(_NDArray_2 == NDArray.scalar(Value.int(Int(0))))]\n",
"_OptionalIntOrTuple_1 = OptionalIntOrTuple.some(IntOrTuple.int(Int(0)))\n",
"_NDArray_4[_IndexKey_1] = sum(_NDArray_5, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_5.shape[Int(0)]))\n",
"_IndexKey_2 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(1))) + _MultiAxisIndexKey_1)\n",
"_NDArray_6 = _NDArray_1[ndarray_index(_NDArray_2 == NDArray.scalar(Value.int(Int(1))))]\n",
"_NDArray_4[_IndexKey_2] = sum(_NDArray_6, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_6.shape[Int(0)]))\n",
"_NDArray_7 = concat(TupleNDArray(_NDArray_5 - _NDArray_4[_IndexKey_1]) + TupleNDArray(_NDArray_6 - _NDArray_4[_IndexKey_2]), OptionalInt.some(Int(0)))\n",
"_NDArray_8 = square(_NDArray_7 - expand_dims(sum(_NDArray_7, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_7.shape[Int(0)]))))\n",
"_NDArray_9 = sqrt(sum(_NDArray_8, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_8.shape[Int(0)])))\n",
"_NDArray_10 = copy(_NDArray_9)\n",
"_NDArray_10[ndarray_index(_NDArray_9 == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(Value.float(Float(1.0)))\n",
"_TupleNDArray_1 = svd(sqrt(NDArray.scalar(Value.float(Float.rational(Rational(1, 999998))))) * (_NDArray_7 / _NDArray_10), FALSE)\n",
"_Slice_1 = Slice(OptionalInt.none, OptionalInt.some(sum(astype(_TupleNDArray_1[Int(1)] > NDArray.scalar(Value.float(Float(0.0001))), DType.int32)).to_value().to_int))\n",
"_NDArray_11 = (_TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(_Slice_1)) + _MultiAxisIndexKey_1)] / _NDArray_10).T / _TupleNDArray_1[\n",
" Int(1)\n",
"][IndexKey.slice(_Slice_1)]\n",
"_TupleNDArray_2 = svd(\n",
" (sqrt((NDArray.scalar(Value.int(Int(1000000))) * _NDArray_3) * NDArray.scalar(Value.float(Float(1.0)))) * (_NDArray_4 - (_NDArray_3 @ _NDArray_4)).T).T @ _NDArray_11, FALSE\n",
")\n",
"(\n",
" (_NDArray_1 - (_NDArray_3 @ _NDArray_4))\n",
" @ (\n",
" _NDArray_11\n",
" @ _TupleNDArray_2[Int(2)].T[\n",
" IndexKey.multi_axis(\n",
" _MultiAxisIndexKey_1\n",
" + MultiAxisIndexKey(\n",
" MultiAxisIndexKeyItem.slice(\n",
" Slice(\n",
" OptionalInt.none,\n",
" OptionalInt.some(\n",
" sum(astype(_TupleNDArray_2[Int(1)] > (NDArray.scalar(Value.float(Float(0.0001))) * _TupleNDArray_2[Int(1)][IndexKey.int(Int(0))]), DType.int32))\n",
" .to_value()\n",
" .to_int\n",
" ),\n",
" )\n",
" )\n",
" )\n",
" )\n",
" ]\n",
" )\n",
")[IndexKey.multi_axis(_MultiAxisIndexKey_1 + MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(1))))))]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"optimized_fn.expr"
]
},
{
"cell_type": "markdown",
"id": "2904e3d2",
"metadata": {},
"source": [
"And the generated code:\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "c606c1bf",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"def __fn(X, y):\n",
" assert X.dtype == np.dtype(np.float64)\n",
" assert X.shape == (1000000, 20,)\n",
" assert np.all(np.isfinite(X))\n",
" assert y.dtype == np.dtype(np.int64)\n",
" assert y.shape == (1000000,)\n",
" assert set(np.unique(y)) == set((0, 1,))\n",
" _0 = y == np.array(0)\n",
" _1 = np.sum(_0)\n",
" _2 = y == np.array(1)\n",
" _3 = np.sum(_2)\n",
" _4 = np.array((_1, _3,)).astype(np.dtype(np.float64))\n",
" _5 = _4 / np.array(1000000.0)\n",
" _6 = np.zeros((2, 20,), dtype=np.dtype(np.float64))\n",
" _7 = np.sum(X[_0], axis=0)\n",
" _8 = _7 / np.array(X[_0].shape[0])\n",
" _6[0, :] = _8\n",
" _9 = np.sum(X[_2], axis=0)\n",
" _10 = _9 / np.array(X[_2].shape[0])\n",
" _6[1, :] = _10\n",
" _11 = _5 @ _6\n",
" _12 = X - _11\n",
" _13 = np.sqrt(np.array(float(1 / 999998)))\n",
" _14 = X[_0] - _6[0, :]\n",
" _15 = X[_2] - _6[1, :]\n",
" _16 = np.concatenate((_14, _15,), axis=0)\n",
" _17 = np.sum(_16, axis=0)\n",
" _18 = _17 / np.array(_16.shape[0])\n",
" _19 = np.expand_dims(_18, 0)\n",
" _20 = _16 - _19\n",
" _21 = np.square(_20)\n",
" _22 = np.sum(_21, axis=0)\n",
" _23 = _22 / np.array(_21.shape[0])\n",
" _24 = np.sqrt(_23)\n",
" _25 = _24 == np.array(0)\n",
" _24[_25] = np.array(1.0)\n",
" _26 = _16 / _24\n",
" _27 = _13 * _26\n",
" _28 = np.linalg.svd(_27, full_matrices=False)\n",
" _29 = _28[1] > np.array(0.0001)\n",
" _30 = _29.astype(np.dtype(np.int32))\n",
" _31 = np.sum(_30)\n",
" _32 = _28[2][:_31, :] / _24\n",
" _33 = _32.T / _28[1][:_31]\n",
" _34 = np.array(1000000) * _5\n",
" _35 = _34 * np.array(1.0)\n",
" _36 = np.sqrt(_35)\n",
" _37 = _6 - _11\n",
" _38 = _36 * _37.T\n",
" _39 = _38.T @ _33\n",
" _40 = np.linalg.svd(_39, full_matrices=False)\n",
" _41 = np.array(0.0001) * _40[1][0]\n",
" _42 = _40[1] > _41\n",
" _43 = _42.astype(np.dtype(np.int32))\n",
" _44 = np.sum(_43)\n",
" _45 = _33 @ _40[2].T[:, :_44]\n",
" _46 = _12 @ _45\n",
" return _46[:, :1]\n",
"\n"
]
}
],
"source": [
"import inspect\n",
"\n",
"print(inspect.getsource(optimized_fn))"
]
},
{
"cell_type": "markdown",
"id": "d6862c1b",
"metadata": {},
"source": [
"As well as the e-graph:\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "f64148d6",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"