_NDArray_1 = NDArray.var("X")\n",
+ ".output_html .il { color: #666 } /* Literal.Number.Integer.Long */_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_shape(_NDArray_1, TupleInt.from_vec(Vec[Int](Int(150), Int(4))))\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 = asarray(reshape(asarray(_NDArray_2), TupleInt(Int(-1))))\n",
- "_NDArray_4 = astype(unique_counts(_NDArray_3)[Int(1)], asarray(_NDArray_1).dtype) / NDArray.scalar(Value.float(Float(1000000.0)))\n",
- "_NDArray_5 = zeros(\n",
- " TupleInt(unique_inverse(_NDArray_3)[Int(0)].shape[Int(0)]) + TupleInt(asarray(_NDArray_1).shape[Int(1)]),\n",
+ "assume_shape(_NDArray_2, TupleInt.from_vec(Vec[Int](Int(150))))\n",
+ "_TupleValue_1 = TupleValue.from_vec(Vec[Value](Value.int(Int(0)), Value.int(Int(1)), Value.int(Int(2))))\n",
+ "assume_value_one_of(_NDArray_2, _TupleValue_1)\n",
+ "_NDArray_3 = zeros(\n",
+ " TupleInt.from_vec(Vec[Int](NDArray.vector(_TupleValue_1).shape[Int(0)], asarray(_NDArray_1).shape[Int(1)])),\n",
" OptionalDType.some(asarray(_NDArray_1).dtype),\n",
" OptionalDevice.some(asarray(_NDArray_1).device),\n",
")\n",
- "_MultiAxisIndexKey_1 = MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice()))\n",
- "_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(0))) + _MultiAxisIndexKey_1)\n",
+ "_MultiAxisIndexKeyItem_1 = MultiAxisIndexKeyItem.slice(Slice())\n",
+ "_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.int(Int(0)), _MultiAxisIndexKeyItem_1)))\n",
+ "_IndexKey_2 = IndexKey.ndarray(unique_inverse(_NDArray_2)[Int(1)] == NDArray.scalar(Value.int(Int(0))))\n",
"_OptionalIntOrTuple_1 = OptionalIntOrTuple.some(IntOrTuple.int(Int(0)))\n",
- "_NDArray_5[_IndexKey_1] = mean(asarray(_NDArray_1)[ndarray_index(unique_inverse(_NDArray_3)[Int(1)] == NDArray.scalar(Value.int(Int(0))))], _OptionalIntOrTuple_1)\n",
- "_IndexKey_2 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(1))) + _MultiAxisIndexKey_1)\n",
- "_NDArray_5[_IndexKey_2] = mean(asarray(_NDArray_1)[ndarray_index(unique_inverse(_NDArray_3)[Int(1)] == NDArray.scalar(Value.int(Int(1))))], _OptionalIntOrTuple_1)\n",
- "_NDArray_6 = unique_values(concat(TupleNDArray(unique_values(asarray(_NDArray_3)))))\n",
- "_NDArray_7 = concat(\n",
- " TupleNDArray(asarray(_NDArray_1)[ndarray_index(_NDArray_3 == _NDArray_6[IndexKey.int(Int(0))])] - _NDArray_5[_IndexKey_1])\n",
- " + TupleNDArray(asarray(_NDArray_1)[ndarray_index(_NDArray_3 == _NDArray_6[IndexKey.int(Int(1))])] - _NDArray_5[_IndexKey_2]),\n",
+ "_NDArray_3[_IndexKey_1] = mean(asarray(_NDArray_1)[_IndexKey_2], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_3 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.int(Int(1)), _MultiAxisIndexKeyItem_1)))\n",
+ "_IndexKey_4 = IndexKey.ndarray(unique_inverse(_NDArray_2)[Int(1)] == NDArray.scalar(Value.int(Int(1))))\n",
+ "_NDArray_3[_IndexKey_3] = mean(asarray(_NDArray_1)[_IndexKey_4], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_5 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.int(Int(2)), _MultiAxisIndexKeyItem_1)))\n",
+ "_IndexKey_6 = IndexKey.ndarray(unique_inverse(_NDArray_2)[Int(1)] == NDArray.scalar(Value.int(Int(2))))\n",
+ "_NDArray_3[_IndexKey_5] = mean(asarray(_NDArray_1)[_IndexKey_6], _OptionalIntOrTuple_1)\n",
+ "_NDArray_4 = zeros(TupleInt.from_vec(Vec[Int](Int(3), Int(4))), OptionalDType.some(DType.float64), OptionalDevice.some(_NDArray_1.device))\n",
+ "_IndexKey_7 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(0)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_4[_IndexKey_7] = mean(_NDArray_1[_IndexKey_2], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_8 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(1)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_4[_IndexKey_8] = mean(_NDArray_1[_IndexKey_4], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_9 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(2)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_4[_IndexKey_9] = mean(_NDArray_1[_IndexKey_6], _OptionalIntOrTuple_1)\n",
+ "_NDArray_5 = concat(\n",
+ " TupleNDArray.from_vec(\n",
+ " Vec[NDArray](\n",
+ " _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(0))))] - _NDArray_4[_IndexKey_7],\n",
+ " _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(1))))] - _NDArray_4[_IndexKey_8],\n",
+ " _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(2))))] - _NDArray_4[_IndexKey_9],\n",
+ " )\n",
+ " ),\n",
" OptionalInt.some(Int(0)),\n",
")\n",
- "_NDArray_8 = std(_NDArray_7, _OptionalIntOrTuple_1)\n",
- "_NDArray_8[ndarray_index(std(_NDArray_7, _OptionalIntOrTuple_1) == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(Value.float(Float(1.0)))\n",
+ "_NDArray_6 = std(_NDArray_5, _OptionalIntOrTuple_1)\n",
+ "_NDArray_6[IndexKey.ndarray(std(_NDArray_5, _OptionalIntOrTuple_1) == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(\n",
+ " Value.float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("1"))))\n",
+ ")\n",
"_TupleNDArray_1 = svd(\n",
- " sqrt(asarray(NDArray.scalar(Value.float(Float(1.0) / Float.from_int(asarray(_NDArray_1).shape[Int(0)] - _NDArray_6.shape[Int(0)]))))) * (_NDArray_7 / _NDArray_8), FALSE\n",
+ " sqrt(\n",
+ " asarray(\n",
+ " NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("147"))))),\n",
+ " OptionalDType.some(DType.float64),\n",
+ " OptionalBool.none,\n",
+ " OptionalDevice.some(_NDArray_1.device),\n",
+ " )\n",
+ " )\n",
+ " * (_NDArray_5 / _NDArray_6),\n",
+ " Boolean(False),\n",
")\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_9 = (_TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(_Slice_1)) + _MultiAxisIndexKey_1)] / _NDArray_8).T / _TupleNDArray_1[\n",
- " Int(1)\n",
- "][IndexKey.slice(_Slice_1)]\n",
+ "_NDArray_7 = asarray(reshape(asarray(_NDArray_2), TupleInt.from_vec(Vec[Int](Int(-1)))))\n",
+ "_NDArray_8 = unique_values(concat(TupleNDArray.from_vec(Vec[NDArray](unique_values(asarray(_NDArray_7))))))\n",
+ "_NDArray_9 = std(\n",
+ " concat(\n",
+ " TupleNDArray.from_vec(\n",
+ " Vec[NDArray](\n",
+ " asarray(_NDArray_1)[IndexKey.ndarray(_NDArray_7 == _NDArray_8[IndexKey.int(Int(0))])] - _NDArray_3[_IndexKey_1],\n",
+ " asarray(_NDArray_1)[IndexKey.ndarray(_NDArray_7 == _NDArray_8[IndexKey.int(Int(1))])] - _NDArray_3[_IndexKey_3],\n",
+ " asarray(_NDArray_1)[IndexKey.ndarray(_NDArray_7 == _NDArray_8[IndexKey.int(Int(2))])] - _NDArray_3[_IndexKey_5],\n",
+ " )\n",
+ " ),\n",
+ " OptionalInt.some(Int(0)),\n",
+ " ),\n",
+ " _OptionalIntOrTuple_1,\n",
+ ")\n",
+ "_NDArray_10 = copy(_NDArray_9)\n",
+ "_NDArray_10[IndexKey.ndarray(_NDArray_9 == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(Value.float(Float(1.0)))\n",
+ "_NDArray_11 = astype(unique_counts(_NDArray_2)[Int(1)], DType.float64) / NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string("150"), BigInt.from_string("1")))))\n",
"_TupleNDArray_2 = svd(\n",
" (\n",
- " sqrt(\n",
- " (NDArray.scalar(Value.int(asarray(_NDArray_1).shape[Int(0)])) * _NDArray_4)\n",
- " * NDArray.scalar(Value.float(Float(1.0) / Float.from_int(_NDArray_6.shape[Int(0)] - Int(1))))\n",
- " )\n",
- " * (_NDArray_5 - (_NDArray_4 @ _NDArray_5)).T\n",
+ " sqrt((NDArray.scalar(Value.int(Int(150))) * _NDArray_11) * NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("2"))))))\n",
+ " * (_NDArray_4 - (_NDArray_11 @ _NDArray_4)).T\n",
" ).T\n",
- " @ _NDArray_9,\n",
- " FALSE,\n",
+ " @ (\n",
+ " (\n",
+ " _TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.slice(_Slice_1), _MultiAxisIndexKeyItem_1)))]\n",
+ " / _NDArray_6\n",
+ " ).T\n",
+ " / _TupleNDArray_1[Int(1)][IndexKey.slice(_Slice_1)]\n",
+ " ),\n",
+ " Boolean(False),\n",
")\n",
"(\n",
- " (asarray(_NDArray_1) - (_NDArray_4 @ _NDArray_5))\n",
+ " (asarray(_NDArray_1) - ((astype(unique_counts(_NDArray_2)[Int(1)], asarray(_NDArray_1).dtype) / NDArray.scalar(Value.float(Float(150.0)))) @ _NDArray_3))\n",
" @ (\n",
- " _NDArray_9\n",
+ " (\n",
+ " (_TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.slice(_Slice_1), _MultiAxisIndexKeyItem_1)))] / _NDArray_10).T\n",
+ " / _TupleNDArray_1[Int(1)][IndexKey.slice(_Slice_1)]\n",
+ " )\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",
+ " MultiAxisIndexKey.from_vec(\n",
+ " Vec(\n",
+ " _MultiAxisIndexKeyItem_1,\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",
" )\n",
- ")[IndexKey.multi_axis(_MultiAxisIndexKey_1 + MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(_NDArray_6.shape[Int(0)] - Int(1))))))]\n",
+ ")[IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(_MultiAxisIndexKeyItem_1, MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(2)))))))]\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{}shape}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{150}\\PY{p}{)}\\PY{p}{,} \\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\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}{asarray}\\PY{p}{(}\\PY{n}{reshape}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{)}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{o}{\\PYZhy{}}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
- "\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{=} \\PY{n}{astype}\\PY{p}{(}\\PY{n}{unique\\PYZus{}counts}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{]}\\PY{p}{,} \\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{dtype}\\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{}5} \\PY{o}{=} \\PY{n}{zeros}\\PY{p}{(}\n",
- " \\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{unique\\PYZus{}inverse}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)} \\PY{o}{+} \\PY{n}{TupleInt}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)}\\PY{p}{,}\n",
+ "\\PY{n}{assume\\PYZus{}shape}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{150}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}TupleValue\\PYZus{}1} \\PY{o}{=} \\PY{n}{TupleValue}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Value}\\PY{p}{]}\\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{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{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\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}{\\PYZus{}TupleValue\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}3} \\PY{o}{=} \\PY{n}{zeros}\\PY{p}{(}\n",
+ " \\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{vector}\\PY{p}{(}\\PY{n}{\\PYZus{}TupleValue\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]}\\PY{p}{,} \\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{shape}\\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}{OptionalDType}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{dtype}\\PY{p}{)}\\PY{p}{,}\n",
" \\PY{n}{OptionalDevice}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{device}\\PY{p}{)}\\PY{p}{,}\n",
"\\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{}MultiAxisIndexKeyItem\\PYZus{}1} \\PY{o}{=} \\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{Slice}\\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{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\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{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{unique\\PYZus{}inverse}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\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}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\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{}5}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}1}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{unique\\PYZus{}inverse}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\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}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\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{}5}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{unique\\PYZus{}inverse}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\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}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
- "\\PY{n}{\\PYZus{}NDArray\\PYZus{}6} \\PY{o}{=} \\PY{n}{unique\\PYZus{}values}\\PY{p}{(}\\PY{n}{concat}\\PY{p}{(}\\PY{n}{TupleNDArray}\\PY{p}{(}\\PY{n}{unique\\PYZus{}values}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
- "\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{=} \\PY{n}{concat}\\PY{p}{(}\n",
- " \\PY{n}{TupleNDArray}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3} \\PY{o}{==} \\PY{n}{\\PYZus{}NDArray\\PYZus{}6}\\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{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}5}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}1}\\PY{p}{]}\\PY{p}{)}\n",
- " \\PY{o}{+} \\PY{n}{TupleNDArray}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}3} \\PY{o}{==} \\PY{n}{\\PYZus{}NDArray\\PYZus{}6}\\PY{p}{[}\\PY{n}{IndexKey}\\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}{]} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}5}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2}\\PY{p}{]}\\PY{p}{)}\\PY{p}{,}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}1}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}IndexKey\\PYZus{}3} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\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{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}IndexKey\\PYZus{}4} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{unique\\PYZus{}inverse}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\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}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}3}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}4}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}IndexKey\\PYZus{}5} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,} \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}IndexKey\\PYZus{}6} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{unique\\PYZus{}inverse}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\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}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}5}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}6}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{=} \\PY{n}{zeros}\\PY{p}{(}\\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{3}\\PY{p}{)}\\PY{p}{,} \\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\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{}IndexKey\\PYZus{}7} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\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{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}7}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}2}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}IndexKey\\PYZus{}8} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\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{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}8}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}4}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}IndexKey\\PYZus{}9} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,} \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}9}\\PY{p}{]} \\PY{o}{=} \\PY{n}{mean}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}6}\\PY{p}{]}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}5} \\PY{o}{=} \\PY{n}{concat}\\PY{p}{(}\n",
+ " \\PY{n}{TupleNDArray}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\n",
+ " \\PY{n}{Vec}\\PY{p}{[}\\PY{n}{NDArray}\\PY{p}{]}\\PY{p}{(}\n",
+ " \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\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}{]} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}7}\\PY{p}{]}\\PY{p}{,}\n",
+ " \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\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}{]} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}8}\\PY{p}{]}\\PY{p}{,}\n",
+ " \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\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}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}9}\\PY{p}{]}\\PY{p}{,}\n",
+ " \\PY{p}{)}\n",
+ " \\PY{p}{)}\\PY{p}{,}\n",
" \\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{p}{)}\n",
- "\\PY{n}{\\PYZus{}NDArray\\PYZus{}8} \\PY{o}{=} \\PY{n}{std}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
- "\\PY{n}{\\PYZus{}NDArray\\PYZus{}8}\\PY{p}{[}\\PY{n}{ndarray\\PYZus{}index}\\PY{p}{(}\\PY{n}{std}\\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}{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{}NDArray\\PYZus{}6} \\PY{o}{=} \\PY{n}{std}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}5}\\PY{p}{,} \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}6}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{std}\\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}{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}{(}\n",
+ " \\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{o}{.}\\PY{n}{rational}\\PY{p}{(}\\PY{n}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1} \\PY{o}{=} \\PY{n}{svd}\\PY{p}{(}\n",
- " \\PY{n}{sqrt}\\PY{p}{(}\\PY{n}{asarray}\\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}{1.0}\\PY{p}{)} \\PY{o}{/} \\PY{n}{Float}\\PY{o}{.}\\PY{n}{from\\PYZus{}int}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{shape}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{]} \\PY{o}{\\PYZhy{}} \\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}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{*} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}8}\\PY{p}{)}\\PY{p}{,} \\PY{n}{FALSE}\n",
+ " \\PY{n}{sqrt}\\PY{p}{(}\n",
+ " \\PY{n}{asarray}\\PY{p}{(}\n",
+ " \\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}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{147}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{OptionalDType}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{DType}\\PY{o}{.}\\PY{n}{float64}\\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{OptionalBool}\\PY{o}{.}\\PY{n}{none}\\PY{p}{,}\n",
+ " \\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{p}{)}\n",
+ " \\PY{p}{)}\n",
+ " \\PY{o}{*} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}5} \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}6}\\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{Boolean}\\PY{p}{(}\\PY{k+kc}{False}\\PY{p}{)}\\PY{p}{,}\n",
"\\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{}9} \\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{}8}\\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{}NDArray\\PYZus{}7} \\PY{o}{=} \\PY{n}{asarray}\\PY{p}{(}\\PY{n}{reshape}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{)}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{o}{\\PYZhy{}}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}8} \\PY{o}{=} \\PY{n}{unique\\PYZus{}values}\\PY{p}{(}\\PY{n}{concat}\\PY{p}{(}\\PY{n}{TupleNDArray}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{NDArray}\\PY{p}{]}\\PY{p}{(}\\PY{n}{unique\\PYZus{}values}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}9} \\PY{o}{=} \\PY{n}{std}\\PY{p}{(}\n",
+ " \\PY{n}{concat}\\PY{p}{(}\n",
+ " \\PY{n}{TupleNDArray}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\n",
+ " \\PY{n}{Vec}\\PY{p}{[}\\PY{n}{NDArray}\\PY{p}{]}\\PY{p}{(}\n",
+ " \\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{==} \\PY{n}{\\PYZus{}NDArray\\PYZus{}8}\\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{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}1}\\PY{p}{]}\\PY{p}{,}\n",
+ " \\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{==} \\PY{n}{\\PYZus{}NDArray\\PYZus{}8}\\PY{p}{[}\\PY{n}{IndexKey}\\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}{]} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}3}\\PY{p}{]}\\PY{p}{,}\n",
+ " \\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{==} \\PY{n}{\\PYZus{}NDArray\\PYZus{}8}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\\PY{p}{)}\\PY{p}{]} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}5}\\PY{p}{]}\\PY{p}{,}\n",
+ " \\PY{p}{)}\n",
+ " \\PY{p}{)}\\PY{p}{,}\n",
+ " \\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{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{\\PYZus{}OptionalIntOrTuple\\PYZus{}1}\\PY{p}{,}\n",
+ "\\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}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\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{}NDArray\\PYZus{}11} \\PY{o}{=} \\PY{n}{astype}\\PY{p}{(}\\PY{n}{unique\\PYZus{}counts}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{]}\\PY{p}{,} \\PY{n}{DType}\\PY{o}{.}\\PY{n}{float64}\\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{o}{.}\\PY{n}{rational}\\PY{p}{(}\\PY{n}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{150}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
"\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}2} \\PY{o}{=} \\PY{n}{svd}\\PY{p}{(}\n",
" \\PY{p}{(}\n",
- " \\PY{n}{sqrt}\\PY{p}{(}\n",
- " \\PY{p}{(}\\PY{n}{NDArray}\\PY{o}{.}\\PY{n}{scalar}\\PY{p}{(}\\PY{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\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{o}{*} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{)}\n",
- " \\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{o}{/} \\PY{n}{Float}\\PY{o}{.}\\PY{n}{from\\PYZus{}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{o}{\\PYZhy{}} \\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
- " \\PY{p}{)}\n",
- " \\PY{o}{*} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}5} \\PY{o}{\\PYZhy{}} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}5}\\PY{p}{)}\\PY{p}{)}\\PY{o}{.}\\PY{n}{T}\n",
+ " \\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}{150}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{*} \\PY{n}{\\PYZus{}NDArray\\PYZus{}11}\\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{o}{.}\\PY{n}{rational}\\PY{p}{(}\\PY{n}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{2}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ " \\PY{o}{*} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{\\PYZhy{}} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}11} \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{)}\\PY{p}{)}\\PY{o}{.}\\PY{n}{T}\n",
" \\PY{p}{)}\\PY{o}{.}\\PY{n}{T}\n",
- " \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}9}\\PY{p}{,}\n",
- " \\PY{n}{FALSE}\\PY{p}{,}\n",
+ " \\PY{o}{@} \\PY{p}{(}\n",
+ " \\PY{p}{(}\n",
+ " \\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{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{\\PYZus{}Slice\\PYZus{}1}\\PY{p}{)}\\PY{p}{,} \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\n",
+ " \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}6}\n",
+ " \\PY{p}{)}\\PY{o}{.}\\PY{n}{T}\n",
+ " \\PY{o}{/} \\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1}\\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}{slice}\\PY{p}{(}\\PY{n}{\\PYZus{}Slice\\PYZus{}1}\\PY{p}{)}\\PY{p}{]}\n",
+ " \\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{Boolean}\\PY{p}{(}\\PY{k+kc}{False}\\PY{p}{)}\\PY{p}{,}\n",
"\\PY{p}{)}\n",
"\\PY{p}{(}\n",
- " \\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)} \\PY{o}{\\PYZhy{}} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}5}\\PY{p}{)}\\PY{p}{)}\n",
+ " \\PY{p}{(}\\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)} \\PY{o}{\\PYZhy{}} \\PY{p}{(}\\PY{p}{(}\\PY{n}{astype}\\PY{p}{(}\\PY{n}{unique\\PYZus{}counts}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{)}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{1}\\PY{p}{)}\\PY{p}{]}\\PY{p}{,} \\PY{n}{asarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{)}\\PY{o}{.}\\PY{n}{dtype}\\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}{150.0}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)} \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}3}\\PY{p}{)}\\PY{p}{)}\n",
" \\PY{o}{@} \\PY{p}{(}\n",
- " \\PY{n}{\\PYZus{}NDArray\\PYZus{}9}\n",
+ " \\PY{p}{(}\n",
+ " \\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{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{\\PYZus{}Slice\\PYZus{}1}\\PY{p}{)}\\PY{p}{,} \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]} \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}10}\\PY{p}{)}\\PY{o}{.}\\PY{n}{T}\n",
+ " \\PY{o}{/} \\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1}\\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}{slice}\\PY{p}{(}\\PY{n}{\\PYZus{}Slice\\PYZus{}1}\\PY{p}{)}\\PY{p}{]}\n",
+ " \\PY{p}{)}\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{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\n",
+ " \\PY{n}{Vec}\\PY{p}{(}\n",
+ " \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\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}{)}\\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}{\\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{o}{\\PYZhy{}} \\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",
+ "\\PY{p}{)}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{(}\\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\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}{2}\\PY{p}{)}\\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_shape(_NDArray_1, TupleInt.from_vec(Vec[Int](Int(150), Int(4))))\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 = asarray(reshape(asarray(_NDArray_2), TupleInt(Int(-1))))\n",
- "_NDArray_4 = astype(unique_counts(_NDArray_3)[Int(1)], asarray(_NDArray_1).dtype) / NDArray.scalar(Value.float(Float(1000000.0)))\n",
- "_NDArray_5 = zeros(\n",
- " TupleInt(unique_inverse(_NDArray_3)[Int(0)].shape[Int(0)]) + TupleInt(asarray(_NDArray_1).shape[Int(1)]),\n",
+ "assume_shape(_NDArray_2, TupleInt.from_vec(Vec[Int](Int(150))))\n",
+ "_TupleValue_1 = TupleValue.from_vec(Vec[Value](Value.int(Int(0)), Value.int(Int(1)), Value.int(Int(2))))\n",
+ "assume_value_one_of(_NDArray_2, _TupleValue_1)\n",
+ "_NDArray_3 = zeros(\n",
+ " TupleInt.from_vec(Vec[Int](NDArray.vector(_TupleValue_1).shape[Int(0)], asarray(_NDArray_1).shape[Int(1)])),\n",
" OptionalDType.some(asarray(_NDArray_1).dtype),\n",
" OptionalDevice.some(asarray(_NDArray_1).device),\n",
")\n",
- "_MultiAxisIndexKey_1 = MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice()))\n",
- "_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(0))) + _MultiAxisIndexKey_1)\n",
+ "_MultiAxisIndexKeyItem_1 = MultiAxisIndexKeyItem.slice(Slice())\n",
+ "_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.int(Int(0)), _MultiAxisIndexKeyItem_1)))\n",
+ "_IndexKey_2 = IndexKey.ndarray(unique_inverse(_NDArray_2)[Int(1)] == NDArray.scalar(Value.int(Int(0))))\n",
"_OptionalIntOrTuple_1 = OptionalIntOrTuple.some(IntOrTuple.int(Int(0)))\n",
- "_NDArray_5[_IndexKey_1] = mean(asarray(_NDArray_1)[ndarray_index(unique_inverse(_NDArray_3)[Int(1)] == NDArray.scalar(Value.int(Int(0))))], _OptionalIntOrTuple_1)\n",
- "_IndexKey_2 = IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.int(Int(1))) + _MultiAxisIndexKey_1)\n",
- "_NDArray_5[_IndexKey_2] = mean(asarray(_NDArray_1)[ndarray_index(unique_inverse(_NDArray_3)[Int(1)] == NDArray.scalar(Value.int(Int(1))))], _OptionalIntOrTuple_1)\n",
- "_NDArray_6 = unique_values(concat(TupleNDArray(unique_values(asarray(_NDArray_3)))))\n",
- "_NDArray_7 = concat(\n",
- " TupleNDArray(asarray(_NDArray_1)[ndarray_index(_NDArray_3 == _NDArray_6[IndexKey.int(Int(0))])] - _NDArray_5[_IndexKey_1])\n",
- " + TupleNDArray(asarray(_NDArray_1)[ndarray_index(_NDArray_3 == _NDArray_6[IndexKey.int(Int(1))])] - _NDArray_5[_IndexKey_2]),\n",
+ "_NDArray_3[_IndexKey_1] = mean(asarray(_NDArray_1)[_IndexKey_2], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_3 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.int(Int(1)), _MultiAxisIndexKeyItem_1)))\n",
+ "_IndexKey_4 = IndexKey.ndarray(unique_inverse(_NDArray_2)[Int(1)] == NDArray.scalar(Value.int(Int(1))))\n",
+ "_NDArray_3[_IndexKey_3] = mean(asarray(_NDArray_1)[_IndexKey_4], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_5 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.int(Int(2)), _MultiAxisIndexKeyItem_1)))\n",
+ "_IndexKey_6 = IndexKey.ndarray(unique_inverse(_NDArray_2)[Int(1)] == NDArray.scalar(Value.int(Int(2))))\n",
+ "_NDArray_3[_IndexKey_5] = mean(asarray(_NDArray_1)[_IndexKey_6], _OptionalIntOrTuple_1)\n",
+ "_NDArray_4 = zeros(TupleInt.from_vec(Vec[Int](Int(3), Int(4))), OptionalDType.some(DType.float64), OptionalDevice.some(_NDArray_1.device))\n",
+ "_IndexKey_7 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(0)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_4[_IndexKey_7] = mean(_NDArray_1[_IndexKey_2], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_8 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(1)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_4[_IndexKey_8] = mean(_NDArray_1[_IndexKey_4], _OptionalIntOrTuple_1)\n",
+ "_IndexKey_9 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(2)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_4[_IndexKey_9] = mean(_NDArray_1[_IndexKey_6], _OptionalIntOrTuple_1)\n",
+ "_NDArray_5 = concat(\n",
+ " TupleNDArray.from_vec(\n",
+ " Vec[NDArray](\n",
+ " _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(0))))] - _NDArray_4[_IndexKey_7],\n",
+ " _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(1))))] - _NDArray_4[_IndexKey_8],\n",
+ " _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(2))))] - _NDArray_4[_IndexKey_9],\n",
+ " )\n",
+ " ),\n",
" OptionalInt.some(Int(0)),\n",
")\n",
- "_NDArray_8 = std(_NDArray_7, _OptionalIntOrTuple_1)\n",
- "_NDArray_8[ndarray_index(std(_NDArray_7, _OptionalIntOrTuple_1) == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(Value.float(Float(1.0)))\n",
+ "_NDArray_6 = std(_NDArray_5, _OptionalIntOrTuple_1)\n",
+ "_NDArray_6[IndexKey.ndarray(std(_NDArray_5, _OptionalIntOrTuple_1) == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(\n",
+ " Value.float(Float.rational(BigRat(BigInt.from_string(\"1\"), BigInt.from_string(\"1\"))))\n",
+ ")\n",
"_TupleNDArray_1 = svd(\n",
- " sqrt(asarray(NDArray.scalar(Value.float(Float(1.0) / Float.from_int(asarray(_NDArray_1).shape[Int(0)] - _NDArray_6.shape[Int(0)]))))) * (_NDArray_7 / _NDArray_8), FALSE\n",
+ " sqrt(\n",
+ " asarray(\n",
+ " NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string(\"1\"), BigInt.from_string(\"147\"))))),\n",
+ " OptionalDType.some(DType.float64),\n",
+ " OptionalBool.none,\n",
+ " OptionalDevice.some(_NDArray_1.device),\n",
+ " )\n",
+ " )\n",
+ " * (_NDArray_5 / _NDArray_6),\n",
+ " Boolean(False),\n",
")\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_9 = (_TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(_Slice_1)) + _MultiAxisIndexKey_1)] / _NDArray_8).T / _TupleNDArray_1[\n",
- " Int(1)\n",
- "][IndexKey.slice(_Slice_1)]\n",
+ "_NDArray_7 = asarray(reshape(asarray(_NDArray_2), TupleInt.from_vec(Vec[Int](Int(-1)))))\n",
+ "_NDArray_8 = unique_values(concat(TupleNDArray.from_vec(Vec[NDArray](unique_values(asarray(_NDArray_7))))))\n",
+ "_NDArray_9 = std(\n",
+ " concat(\n",
+ " TupleNDArray.from_vec(\n",
+ " Vec[NDArray](\n",
+ " asarray(_NDArray_1)[IndexKey.ndarray(_NDArray_7 == _NDArray_8[IndexKey.int(Int(0))])] - _NDArray_3[_IndexKey_1],\n",
+ " asarray(_NDArray_1)[IndexKey.ndarray(_NDArray_7 == _NDArray_8[IndexKey.int(Int(1))])] - _NDArray_3[_IndexKey_3],\n",
+ " asarray(_NDArray_1)[IndexKey.ndarray(_NDArray_7 == _NDArray_8[IndexKey.int(Int(2))])] - _NDArray_3[_IndexKey_5],\n",
+ " )\n",
+ " ),\n",
+ " OptionalInt.some(Int(0)),\n",
+ " ),\n",
+ " _OptionalIntOrTuple_1,\n",
+ ")\n",
+ "_NDArray_10 = copy(_NDArray_9)\n",
+ "_NDArray_10[IndexKey.ndarray(_NDArray_9 == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(Value.float(Float(1.0)))\n",
+ "_NDArray_11 = astype(unique_counts(_NDArray_2)[Int(1)], DType.float64) / NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string(\"150\"), BigInt.from_string(\"1\")))))\n",
"_TupleNDArray_2 = svd(\n",
" (\n",
- " sqrt(\n",
- " (NDArray.scalar(Value.int(asarray(_NDArray_1).shape[Int(0)])) * _NDArray_4)\n",
- " * NDArray.scalar(Value.float(Float(1.0) / Float.from_int(_NDArray_6.shape[Int(0)] - Int(1))))\n",
- " )\n",
- " * (_NDArray_5 - (_NDArray_4 @ _NDArray_5)).T\n",
+ " sqrt((NDArray.scalar(Value.int(Int(150))) * _NDArray_11) * NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string(\"1\"), BigInt.from_string(\"2\"))))))\n",
+ " * (_NDArray_4 - (_NDArray_11 @ _NDArray_4)).T\n",
" ).T\n",
- " @ _NDArray_9,\n",
- " FALSE,\n",
+ " @ (\n",
+ " (\n",
+ " _TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.slice(_Slice_1), _MultiAxisIndexKeyItem_1)))]\n",
+ " / _NDArray_6\n",
+ " ).T\n",
+ " / _TupleNDArray_1[Int(1)][IndexKey.slice(_Slice_1)]\n",
+ " ),\n",
+ " Boolean(False),\n",
")\n",
"(\n",
- " (asarray(_NDArray_1) - (_NDArray_4 @ _NDArray_5))\n",
+ " (asarray(_NDArray_1) - ((astype(unique_counts(_NDArray_2)[Int(1)], asarray(_NDArray_1).dtype) / NDArray.scalar(Value.float(Float(150.0)))) @ _NDArray_3))\n",
" @ (\n",
- " _NDArray_9\n",
+ " (\n",
+ " (_TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(MultiAxisIndexKeyItem.slice(_Slice_1), _MultiAxisIndexKeyItem_1)))] / _NDArray_10).T\n",
+ " / _TupleNDArray_1[Int(1)][IndexKey.slice(_Slice_1)]\n",
+ " )\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",
+ " MultiAxisIndexKey.from_vec(\n",
+ " Vec(\n",
+ " _MultiAxisIndexKeyItem_1,\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",
" )\n",
- ")[IndexKey.multi_axis(_MultiAxisIndexKey_1 + MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(_NDArray_6.shape[Int(0)] - Int(1))))))]"
+ ")[IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec(_MultiAxisIndexKeyItem_1, MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(2)))))))]"
]
},
"metadata": {},
@@ -518,7 +673,7 @@
},
{
"cell_type": "code",
- "execution_count": 4,
+ "execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
@@ -544,7 +699,7 @@
},
{
"cell_type": "code",
- "execution_count": 5,
+ "execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
@@ -555,7 +710,7 @@
},
{
"cell_type": "code",
- "execution_count": 6,
+ "execution_count": 7,
"metadata": {},
"outputs": [
{
@@ -569,9 +724,9 @@
".output_html .hll { background-color: #ffffcc }\n",
".output_html { background: #f8f8f8; }\n",
".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
- ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
+ ".output_html .err { border: 1px solid #F00 } /* Error */\n",
".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
- ".output_html .o { color: #666666 } /* Operator */\n",
+ ".output_html .o { color: #666 } /* Operator */\n",
".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
@@ -588,34 +743,34 @@
".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
- ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
+ ".output_html .gt { color: #04D } /* Generic.Traceback */\n",
".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
- ".output_html .m { color: #666666 } /* Literal.Number */\n",
+ ".output_html .m { color: #666 } /* Literal.Number */\n",
".output_html .s { color: #BA2121 } /* Literal.String */\n",
".output_html .na { color: #687822 } /* Name.Attribute */\n",
".output_html .nb { color: #008000 } /* Name.Builtin */\n",
- ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
- ".output_html .no { color: #880000 } /* Name.Constant */\n",
- ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
+ ".output_html .nc { color: #00F; font-weight: bold } /* Name.Class */\n",
+ ".output_html .no { color: #800 } /* Name.Constant */\n",
+ ".output_html .nd { color: #A2F } /* Name.Decorator */\n",
".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
- ".output_html .nf { color: #0000FF } /* Name.Function */\n",
+ ".output_html .nf { color: #00F } /* Name.Function */\n",
".output_html .nl { color: #767600 } /* Name.Label */\n",
- ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
+ ".output_html .nn { color: #00F; font-weight: bold } /* Name.Namespace */\n",
".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
".output_html .nv { color: #19177C } /* Name.Variable */\n",
- ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
- ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
- ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
- ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
- ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
- ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
- ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
+ ".output_html .ow { color: #A2F; font-weight: bold } /* Operator.Word */\n",
+ ".output_html .w { color: #BBB } /* Text.Whitespace */\n",
+ ".output_html .mb { color: #666 } /* Literal.Number.Bin */\n",
+ ".output_html .mf { color: #666 } /* Literal.Number.Float */\n",
+ ".output_html .mh { color: #666 } /* Literal.Number.Hex */\n",
+ ".output_html .mi { color: #666 } /* Literal.Number.Integer */\n",
+ ".output_html .mo { color: #666 } /* Literal.Number.Oct */\n",
".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
@@ -630,12 +785,12 @@
".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
- ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
+ ".output_html .fm { color: #00F } /* Name.Function.Magic */\n",
".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
- ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */A()[Slice(OptionalInt.none, OptionalInt.some(Int(1)), OptionalInt.some(Int(2)))]\n",
+ ".output_html .il { color: #666 } /* Literal.Number.Integer.Long */A()[Slice(OptionalInt.none, OptionalInt.some(Int(1)), OptionalInt.some(Int(2)))]\n",
" \n"
],
"text/latex": [
@@ -719,7 +874,7 @@
},
{
"cell_type": "code",
- "execution_count": 7,
+ "execution_count": 8,
"metadata": {},
"outputs": [
{
@@ -733,9 +888,9 @@
".output_html .hll { background-color: #ffffcc }\n",
".output_html { background: #f8f8f8; }\n",
".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
- ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
+ ".output_html .err { border: 1px solid #F00 } /* Error */\n",
".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
- ".output_html .o { color: #666666 } /* Operator */\n",
+ ".output_html .o { color: #666 } /* Operator */\n",
".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
@@ -752,34 +907,34 @@
".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
- ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
+ ".output_html .gt { color: #04D } /* Generic.Traceback */\n",
".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
- ".output_html .m { color: #666666 } /* Literal.Number */\n",
+ ".output_html .m { color: #666 } /* Literal.Number */\n",
".output_html .s { color: #BA2121 } /* Literal.String */\n",
".output_html .na { color: #687822 } /* Name.Attribute */\n",
".output_html .nb { color: #008000 } /* Name.Builtin */\n",
- ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
- ".output_html .no { color: #880000 } /* Name.Constant */\n",
- ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
+ ".output_html .nc { color: #00F; font-weight: bold } /* Name.Class */\n",
+ ".output_html .no { color: #800 } /* Name.Constant */\n",
+ ".output_html .nd { color: #A2F } /* Name.Decorator */\n",
".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
- ".output_html .nf { color: #0000FF } /* Name.Function */\n",
+ ".output_html .nf { color: #00F } /* Name.Function */\n",
".output_html .nl { color: #767600 } /* Name.Label */\n",
- ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
+ ".output_html .nn { color: #00F; font-weight: bold } /* Name.Namespace */\n",
".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
".output_html .nv { color: #19177C } /* Name.Variable */\n",
- ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
- ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
- ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
- ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
- ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
- ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
- ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
+ ".output_html .ow { color: #A2F; font-weight: bold } /* Operator.Word */\n",
+ ".output_html .w { color: #BBB } /* Text.Whitespace */\n",
+ ".output_html .mb { color: #666 } /* Literal.Number.Bin */\n",
+ ".output_html .mf { color: #666 } /* Literal.Number.Float */\n",
+ ".output_html .mh { color: #666 } /* Literal.Number.Hex */\n",
+ ".output_html .mi { color: #666 } /* Literal.Number.Integer */\n",
+ ".output_html .mo { color: #666 } /* Literal.Number.Oct */\n",
".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
@@ -794,12 +949,12 @@
".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
- ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
+ ".output_html .fm { color: #00F } /* Name.Function.Magic */\n",
".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
- ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */Boolean(True) & Boolean(True)\n",
+ ".output_html .il { color: #666 } /* Literal.Number.Integer.Long */Boolean(True) & Boolean(True)\n",
" \n"
],
"text/latex": [
@@ -831,7 +986,7 @@
" egraph = EGraph()\n",
" egraph.register(self)\n",
" egraph.run(bool_rewrites.saturate())\n",
- " return egraph.eval(self.bool)\n",
+ " return egraph.extract(self.bool).value\n",
"\n",
"\n",
"x = var(\"x\", Boolean)\n",
@@ -847,7 +1002,7 @@
},
{
"cell_type": "code",
- "execution_count": 8,
+ "execution_count": 9,
"metadata": {},
"outputs": [
{
@@ -913,158 +1068,18 @@
},
{
"cell_type": "code",
- "execution_count": 9,
+ "execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
- "image/svg+xml": [
- ""
- ],
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "b3a8e2d4b3224529aa8edbc731df10e4",
+ "version_major": 2,
+ "version_minor": 1
+ },
"text/plain": [
- ""
+ "VisualizerWidget(egraphs=['{\"nodes\":{\"primitive-i64-0\":{\"op\":\"0\",\"children\":[],\"eclass\":\"i64-0\",\"cost\":1.0,\"su…"
]
},
"metadata": {},
@@ -1082,9 +1097,9 @@
"xs = ListOfInts()\n",
"xs[0] = Int(1)\n",
"\n",
- "egraph = EGraph()\n",
- "egraph.register(xs[0])\n",
- "egraph.display()"
+ "new_egraph = EGraph()\n",
+ "new_egraph.register(xs[0])\n",
+ "new_egraph.display()"
]
},
{
@@ -1125,7 +1140,7 @@
},
{
"cell_type": "code",
- "execution_count": 10,
+ "execution_count": 11,
"metadata": {},
"outputs": [
{
@@ -1139,9 +1154,9 @@
".output_html .hll { background-color: #ffffcc }\n",
".output_html { background: #f8f8f8; }\n",
".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
- ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
+ ".output_html .err { border: 1px solid #F00 } /* Error */\n",
".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
- ".output_html .o { color: #666666 } /* Operator */\n",
+ ".output_html .o { color: #666 } /* Operator */\n",
".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
@@ -1158,34 +1173,34 @@
".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
- ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
+ ".output_html .gt { color: #04D } /* Generic.Traceback */\n",
".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
- ".output_html .m { color: #666666 } /* Literal.Number */\n",
+ ".output_html .m { color: #666 } /* Literal.Number */\n",
".output_html .s { color: #BA2121 } /* Literal.String */\n",
".output_html .na { color: #687822 } /* Name.Attribute */\n",
".output_html .nb { color: #008000 } /* Name.Builtin */\n",
- ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
- ".output_html .no { color: #880000 } /* Name.Constant */\n",
- ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
+ ".output_html .nc { color: #00F; font-weight: bold } /* Name.Class */\n",
+ ".output_html .no { color: #800 } /* Name.Constant */\n",
+ ".output_html .nd { color: #A2F } /* Name.Decorator */\n",
".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
- ".output_html .nf { color: #0000FF } /* Name.Function */\n",
+ ".output_html .nf { color: #00F } /* Name.Function */\n",
".output_html .nl { color: #767600 } /* Name.Label */\n",
- ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
+ ".output_html .nn { color: #00F; font-weight: bold } /* Name.Namespace */\n",
".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
".output_html .nv { color: #19177C } /* Name.Variable */\n",
- ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
- ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
- ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
- ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
- ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
- ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
- ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
+ ".output_html .ow { color: #A2F; font-weight: bold } /* Operator.Word */\n",
+ ".output_html .w { color: #BBB } /* Text.Whitespace */\n",
+ ".output_html .mb { color: #666 } /* Literal.Number.Bin */\n",
+ ".output_html .mf { color: #666 } /* Literal.Number.Float */\n",
+ ".output_html .mh { color: #666 } /* Literal.Number.Hex */\n",
+ ".output_html .mi { color: #666 } /* Literal.Number.Integer */\n",
+ ".output_html .mo { color: #666 } /* Literal.Number.Oct */\n",
".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
@@ -1200,189 +1215,303 @@
".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
- ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
+ ".output_html .fm { color: #00F } /* Name.Function.Magic */\n",
".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
- ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */_NDArray_1 = NDArray.var("X")\n",
+ ".output_html .il { color: #666 } /* Literal.Number.Integer.Long */_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_shape(_NDArray_1, TupleInt.from_vec(Vec[Int](Int(150), Int(4))))\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",
+ "assume_shape(_NDArray_2, TupleInt.from_vec(Vec[Int](Int(150))))\n",
+ "assume_value_one_of(_NDArray_2, TupleValue.from_vec(Vec[Value](Value.int(Int(0)), Value.int(Int(1)), Value.int(Int(2)))))\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",
+ " NDArray.vector(\n",
+ " TupleValue.from_vec(\n",
+ " Vec[Value](\n",
+ " sum(_NDArray_2 == NDArray.scalar(Value.int(Int(0)))).to_value(),\n",
+ " sum(_NDArray_2 == NDArray.scalar(Value.int(Int(1)))).to_value(),\n",
+ " sum(_NDArray_2 == NDArray.scalar(Value.int(Int(2)))).to_value(),\n",
+ " )\n",
+ " )\n",
+ " ),\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",
+ ") / NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string("150"), BigInt.from_string("1")))))\n",
+ "_NDArray_4 = zeros(TupleInt.from_vec(Vec[Int](Int(3), Int(4))), OptionalDType.some(DType.float64), OptionalDevice.some(_NDArray_1.device))\n",
+ "_MultiAxisIndexKeyItem_1 = MultiAxisIndexKeyItem.slice(Slice())\n",
+ "_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(0)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_5 = _NDArray_1[IndexKey.ndarray(_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",
+ "_IndexKey_2 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(1)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_6 = _NDArray_1[IndexKey.ndarray(_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",
+ "_IndexKey_3 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(2)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_7 = _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(2))))]\n",
+ "_NDArray_4[_IndexKey_3] = sum(_NDArray_7, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_7.shape[Int(0)]))\n",
+ "_NDArray_8 = concat(\n",
+ " TupleNDArray.from_vec(Vec[NDArray](_NDArray_5 - _NDArray_4[_IndexKey_1], _NDArray_6 - _NDArray_4[_IndexKey_2], _NDArray_7 - _NDArray_4[_IndexKey_3])), OptionalInt.some(Int(0))\n",
+ ")\n",
+ "_NDArray_9 = square(_NDArray_8 - expand_dims(sum(_NDArray_8, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_8.shape[Int(0)]))))\n",
+ "_NDArray_10 = sqrt(sum(_NDArray_9, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_9.shape[Int(0)])))\n",
+ "_NDArray_11 = copy(_NDArray_10)\n",
+ "_NDArray_11[IndexKey.ndarray(_NDArray_10 == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(\n",
+ " Value.float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("1"))))\n",
+ ")\n",
+ "_TupleNDArray_1 = svd(\n",
+ " sqrt(\n",
+ " asarray(\n",
+ " NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("147"))))),\n",
+ " OptionalDType.some(DType.float64),\n",
+ " OptionalBool.none,\n",
+ " OptionalDevice.some(_NDArray_1.device),\n",
+ " )\n",
+ " )\n",
+ " * (_NDArray_8 / _NDArray_11),\n",
+ " Boolean(False),\n",
+ ")\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",
+ "_NDArray_12 = (\n",
+ " _TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.slice(_Slice_1), _MultiAxisIndexKeyItem_1)))]\n",
+ " / _NDArray_11\n",
+ ").T / _TupleNDArray_1[Int(1)][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",
+ " sqrt((NDArray.scalar(Value.int(Int(150))) * _NDArray_3) * NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string("1"), BigInt.from_string("2"))))))\n",
+ " * (_NDArray_4 - (_NDArray_3 @ _NDArray_4)).T\n",
+ " ).T\n",
+ " @ _NDArray_12,\n",
+ " Boolean(False),\n",
")\n",
"(\n",
" (_NDArray_1 - (_NDArray_3 @ _NDArray_4))\n",
" @ (\n",
- " _NDArray_11\n",
+ " _NDArray_12\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",
+ " MultiAxisIndexKey.from_vec(\n",
+ " Vec[MultiAxisIndexKeyItem](\n",
+ " _MultiAxisIndexKeyItem_1,\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",
" )\n",
- ")[IndexKey.multi_axis(_MultiAxisIndexKey_1 + MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(1))))))]\n",
+ ")[\n",
+ " IndexKey.multi_axis(\n",
+ " MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](_MultiAxisIndexKeyItem_1, MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(2))))))\n",
+ " )\n",
+ "]\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{}shape}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{150}\\PY{p}{)}\\PY{p}{,} \\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\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}{assume\\PYZus{}shape}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}2}\\PY{p}{,} \\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{150}\\PY{p}{)}\\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{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Value}\\PY{p}{]}\\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{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{n}{Value}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\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}{NDArray}\\PY{o}{.}\\PY{n}{vector}\\PY{p}{(}\n",
+ " \\PY{n}{TupleValue}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\n",
+ " \\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Value}\\PY{p}{]}\\PY{p}{(}\n",
+ " \\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}{,}\n",
+ " \\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}{,}\n",
+ " \\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}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{o}{.}\\PY{n}{to\\PYZus{}value}\\PY{p}{(}\\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{p}{)}\n",
+ " \\PY{p}{)}\n",
+ " \\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{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{o}{.}\\PY{n}{rational}\\PY{p}{(}\\PY{n}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{150}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}4} \\PY{o}{=} \\PY{n}{zeros}\\PY{p}{(}\\PY{n}{TupleInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{Int}\\PY{p}{]}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{3}\\PY{p}{)}\\PY{p}{,} \\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\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{}MultiAxisIndexKeyItem\\PYZus{}1} \\PY{o}{=} \\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{Slice}\\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{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\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{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}5} \\PY{o}{=} \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\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{}IndexKey\\PYZus{}2} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\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{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}6} \\PY{o}{=} \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\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{}IndexKey\\PYZus{}3} \\PY{o}{=} \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{int}\\PY{p}{(}\\PY{n}{Int}\\PY{p}{(}\\PY{l+m+mi}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,} \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{=} \\PY{n}{\\PYZus{}NDArray\\PYZus{}1}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\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}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}3}\\PY{p}{]} \\PY{o}{=} \\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}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}8} \\PY{o}{=} \\PY{n}{concat}\\PY{p}{(}\n",
+ " \\PY{n}{TupleNDArray}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{NDArray}\\PY{p}{]}\\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{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{n}{\\PYZus{}NDArray\\PYZus{}7} \\PY{o}{\\PYZhy{}} \\PY{n}{\\PYZus{}NDArray\\PYZus{}4}\\PY{p}{[}\\PY{n}{\\PYZus{}IndexKey\\PYZus{}3}\\PY{p}{]}\\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}{)}\n",
+ "\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}9} \\PY{o}{=} \\PY{n}{square}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}8} \\PY{o}{\\PYZhy{}} \\PY{n}{expand\\PYZus{}dims}\\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}{)}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}10} \\PY{o}{=} \\PY{n}{sqrt}\\PY{p}{(}\\PY{n+nb}{sum}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}9}\\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{}9}\\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{}11} \\PY{o}{=} \\PY{n}{copy}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}10}\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}NDArray\\PYZus{}11}\\PY{p}{[}\\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{ndarray}\\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}10} \\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}{(}\n",
+ " \\PY{n}{Value}\\PY{o}{.}\\PY{n}{float}\\PY{p}{(}\\PY{n}{Float}\\PY{o}{.}\\PY{n}{rational}\\PY{p}{(}\\PY{n}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ "\\PY{p}{)}\n",
+ "\\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1} \\PY{o}{=} \\PY{n}{svd}\\PY{p}{(}\n",
+ " \\PY{n}{sqrt}\\PY{p}{(}\n",
+ " \\PY{n}{asarray}\\PY{p}{(}\n",
+ " \\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}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{147}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{OptionalDType}\\PY{o}{.}\\PY{n}{some}\\PY{p}{(}\\PY{n}{DType}\\PY{o}{.}\\PY{n}{float64}\\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{OptionalBool}\\PY{o}{.}\\PY{n}{none}\\PY{p}{,}\n",
+ " \\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{p}{)}\n",
+ " \\PY{p}{)}\n",
+ " \\PY{o}{*} \\PY{p}{(}\\PY{n}{\\PYZus{}NDArray\\PYZus{}8} \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}11}\\PY{p}{)}\\PY{p}{,}\n",
+ " \\PY{n}{Boolean}\\PY{p}{(}\\PY{k+kc}{False}\\PY{p}{)}\\PY{p}{,}\n",
+ "\\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{}NDArray\\PYZus{}12} \\PY{o}{=} \\PY{p}{(}\n",
+ " \\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{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\PY{p}{(}\\PY{n}{MultiAxisIndexKeyItem}\\PY{o}{.}\\PY{n}{slice}\\PY{p}{(}\\PY{n}{\\PYZus{}Slice\\PYZus{}1}\\PY{p}{)}\\PY{p}{,} \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{]}\n",
+ " \\PY{o}{/} \\PY{n}{\\PYZus{}NDArray\\PYZus{}11}\n",
+ "\\PY{p}{)}\\PY{o}{.}\\PY{n}{T} \\PY{o}{/} \\PY{n}{\\PYZus{}TupleNDArray\\PYZus{}1}\\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}{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{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}{150}\\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{o}{.}\\PY{n}{rational}\\PY{p}{(}\\PY{n}{BigRat}\\PY{p}{(}\\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{1}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{,} \\PY{n}{BigInt}\\PY{o}{.}\\PY{n}{from\\PYZus{}string}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{2}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ " \\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}\n",
+ " \\PY{p}{)}\\PY{o}{.}\\PY{n}{T}\n",
+ " \\PY{o}{@} \\PY{n}{\\PYZus{}NDArray\\PYZus{}12}\\PY{p}{,}\n",
+ " \\PY{n}{Boolean}\\PY{p}{(}\\PY{k+kc}{False}\\PY{p}{)}\\PY{p}{,}\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{n}{\\PYZus{}NDArray\\PYZus{}12}\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{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\n",
+ " \\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\PY{p}{(}\n",
+ " \\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\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}{)}\\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",
+ "\\PY{p}{)}\\PY{p}{[}\n",
+ " \\PY{n}{IndexKey}\\PY{o}{.}\\PY{n}{multi\\PYZus{}axis}\\PY{p}{(}\n",
+ " \\PY{n}{MultiAxisIndexKey}\\PY{o}{.}\\PY{n}{from\\PYZus{}vec}\\PY{p}{(}\\PY{n}{Vec}\\PY{p}{[}\\PY{n}{MultiAxisIndexKeyItem}\\PY{p}{]}\\PY{p}{(}\\PY{n}{\\PYZus{}MultiAxisIndexKeyItem\\PYZus{}1}\\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}{2}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\\PY{p}{)}\n",
+ " \\PY{p}{)}\n",
+ "\\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_shape(_NDArray_1, TupleInt.from_vec(Vec[Int](Int(150), Int(4))))\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",
+ "assume_shape(_NDArray_2, TupleInt.from_vec(Vec[Int](Int(150))))\n",
+ "assume_value_one_of(_NDArray_2, TupleValue.from_vec(Vec[Value](Value.int(Int(0)), Value.int(Int(1)), Value.int(Int(2)))))\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",
+ " NDArray.vector(\n",
+ " TupleValue.from_vec(\n",
+ " Vec[Value](\n",
+ " sum(_NDArray_2 == NDArray.scalar(Value.int(Int(0)))).to_value(),\n",
+ " sum(_NDArray_2 == NDArray.scalar(Value.int(Int(1)))).to_value(),\n",
+ " sum(_NDArray_2 == NDArray.scalar(Value.int(Int(2)))).to_value(),\n",
+ " )\n",
+ " )\n",
+ " ),\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",
+ ") / NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string(\"150\"), BigInt.from_string(\"1\")))))\n",
+ "_NDArray_4 = zeros(TupleInt.from_vec(Vec[Int](Int(3), Int(4))), OptionalDType.some(DType.float64), OptionalDevice.some(_NDArray_1.device))\n",
+ "_MultiAxisIndexKeyItem_1 = MultiAxisIndexKeyItem.slice(Slice())\n",
+ "_IndexKey_1 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(0)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_5 = _NDArray_1[IndexKey.ndarray(_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",
+ "_IndexKey_2 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(1)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_6 = _NDArray_1[IndexKey.ndarray(_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",
+ "_IndexKey_3 = IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.int(Int(2)), _MultiAxisIndexKeyItem_1)))\n",
+ "_NDArray_7 = _NDArray_1[IndexKey.ndarray(_NDArray_2 == NDArray.scalar(Value.int(Int(2))))]\n",
+ "_NDArray_4[_IndexKey_3] = sum(_NDArray_7, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_7.shape[Int(0)]))\n",
+ "_NDArray_8 = concat(\n",
+ " TupleNDArray.from_vec(Vec[NDArray](_NDArray_5 - _NDArray_4[_IndexKey_1], _NDArray_6 - _NDArray_4[_IndexKey_2], _NDArray_7 - _NDArray_4[_IndexKey_3])), OptionalInt.some(Int(0))\n",
+ ")\n",
+ "_NDArray_9 = square(_NDArray_8 - expand_dims(sum(_NDArray_8, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_8.shape[Int(0)]))))\n",
+ "_NDArray_10 = sqrt(sum(_NDArray_9, _OptionalIntOrTuple_1) / NDArray.scalar(Value.int(_NDArray_9.shape[Int(0)])))\n",
+ "_NDArray_11 = copy(_NDArray_10)\n",
+ "_NDArray_11[IndexKey.ndarray(_NDArray_10 == NDArray.scalar(Value.int(Int(0))))] = NDArray.scalar(\n",
+ " Value.float(Float.rational(BigRat(BigInt.from_string(\"1\"), BigInt.from_string(\"1\"))))\n",
+ ")\n",
+ "_TupleNDArray_1 = svd(\n",
+ " sqrt(\n",
+ " asarray(\n",
+ " NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string(\"1\"), BigInt.from_string(\"147\"))))),\n",
+ " OptionalDType.some(DType.float64),\n",
+ " OptionalBool.none,\n",
+ " OptionalDevice.some(_NDArray_1.device),\n",
+ " )\n",
+ " )\n",
+ " * (_NDArray_8 / _NDArray_11),\n",
+ " Boolean(False),\n",
+ ")\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",
+ "_NDArray_12 = (\n",
+ " _TupleNDArray_1[Int(2)][IndexKey.multi_axis(MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](MultiAxisIndexKeyItem.slice(_Slice_1), _MultiAxisIndexKeyItem_1)))]\n",
+ " / _NDArray_11\n",
+ ").T / _TupleNDArray_1[Int(1)][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",
+ " sqrt((NDArray.scalar(Value.int(Int(150))) * _NDArray_3) * NDArray.scalar(Value.float(Float.rational(BigRat(BigInt.from_string(\"1\"), BigInt.from_string(\"2\"))))))\n",
+ " * (_NDArray_4 - (_NDArray_3 @ _NDArray_4)).T\n",
+ " ).T\n",
+ " @ _NDArray_12,\n",
+ " Boolean(False),\n",
")\n",
"(\n",
" (_NDArray_1 - (_NDArray_3 @ _NDArray_4))\n",
" @ (\n",
- " _NDArray_11\n",
+ " _NDArray_12\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",
+ " MultiAxisIndexKey.from_vec(\n",
+ " Vec[MultiAxisIndexKeyItem](\n",
+ " _MultiAxisIndexKeyItem_1,\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",
" )\n",
- ")[IndexKey.multi_axis(_MultiAxisIndexKey_1 + MultiAxisIndexKey(MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(1))))))]"
+ ")[\n",
+ " IndexKey.multi_axis(\n",
+ " MultiAxisIndexKey.from_vec(Vec[MultiAxisIndexKeyItem](_MultiAxisIndexKeyItem_1, MultiAxisIndexKeyItem.slice(Slice(OptionalInt.none, OptionalInt.some(Int(2))))))\n",
+ " )\n",
+ "]"
]
},
"metadata": {},
@@ -1392,7 +1521,10 @@
"source": [
"from egglog.exp.array_api_numba import array_api_numba_schedule\n",
"\n",
- "simplified_res = EGraph().simplify(res, array_api_numba_schedule)\n",
+ "egraph.register(res)\n",
+ "egraph.run(array_api_numba_schedule)\n",
+ "simplified_res = egraph.extract(res)\n",
+ "\n",
"simplified_res"
]
},
@@ -1413,37 +1545,7 @@
"\n",
"Well we showed how we can use eager evaluation to get a result, but what if we don't want to do the computation in egglog, but instead export a program so we can execute that back in Python or in this case feed it to Python?\n",
"\n",
- "Well in this case we have designed a `Program` object which we can use to convert a funtional egglog expression back to imperative Python code:\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 11,
- "metadata": {},
- "outputs": [
- {
- "data": {
- "text/plain": [
- ""
- ]
- },
- "execution_count": 11,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "from egglog.exp.array_api_program_gen import *\n",
- "\n",
- "egraph = EGraph()\n",
- "fn_program = egraph.let(\n",
- " \"fn_program\",\n",
- " ndarray_function_two(simplified_res, NDArray.var(\"X\"), NDArray.var(\"y\")),\n",
- ")\n",
- "egraph.run(array_api_program_gen_schedule)\n",
- "fn = egraph.eval(fn_program.py_object)\n",
- "\n",
- "fn"
+ "Well in this case we have designed a `Program` object which we can use to convert a funtional egglog expression back to imperative Python code:"
]
},
{
@@ -1457,69 +1559,86 @@
"text": [
"def __fn(X, y):\n",
" assert X.dtype == np.dtype(np.float64)\n",
- " assert X.shape == (1000000, 20,)\n",
+ " assert X.shape == (150, 4, )\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",
+ " assert y.shape == (150, )\n",
+ " assert set(np.unique(y)) == set((0, 1, 2, ))\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",
+ " _4 = y == np.array(2)\n",
+ " _5 = np.sum(_4)\n",
+ " _6 = np.array((_1, _3, _5, )).astype(np.dtype(np.float64))\n",
+ " _7 = _6 / np.array(float(150))\n",
+ " _8 = np.zeros((3, 4, ), dtype=np.dtype(np.float64))\n",
+ " _9 = np.sum(X[_0], axis=0)\n",
+ " _10 = _9 / np.array(X[_0].shape[0])\n",
+ " _8[0, :,] = _10\n",
+ " _11 = np.sum(X[_2], axis=0)\n",
+ " _12 = _11 / np.array(X[_2].shape[0])\n",
+ " _8[1, :,] = _12\n",
+ " _13 = np.sum(X[_4], axis=0)\n",
+ " _14 = _13 / np.array(X[_4].shape[0])\n",
+ " _8[2, :,] = _14\n",
+ " _15 = _7 @ _8\n",
+ " _16 = X - _15\n",
+ " _17 = np.sqrt(np.asarray(np.array(float(1 / 147)), np.dtype(np.float64)))\n",
+ " _18 = X[_0] - _8[0, :,]\n",
+ " _19 = X[_2] - _8[1, :,]\n",
+ " _20 = X[_4] - _8[2, :,]\n",
+ " _21 = np.concatenate((_18, _19, _20, ), axis=0)\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",
+ " _24 = np.expand_dims(_23, 0)\n",
+ " _25 = _21 - _24\n",
+ " _26 = np.square(_25)\n",
+ " _27 = np.sum(_26, axis=0)\n",
+ " _28 = _27 / np.array(_26.shape[0])\n",
+ " _29 = np.sqrt(_28)\n",
+ " _30 = _29 == np.array(0)\n",
+ " _29[_30] = np.array(float(1))\n",
+ " _31 = _21 / _29\n",
+ " _32 = _17 * _31\n",
+ " _33 = np.linalg.svd(_32, full_matrices=False)\n",
+ " _34 = _33[1] > np.array(0.0001)\n",
+ " _35 = _34.astype(np.dtype(np.int32))\n",
+ " _36 = np.sum(_35)\n",
+ " _37 = _33[2][:_36, :,] / _29\n",
+ " _38 = _37.T / _33[1][:_36]\n",
+ " _39 = np.array(150) * _7\n",
+ " _40 = _39 * np.array(float(1 / 2))\n",
+ " _41 = np.sqrt(_40)\n",
+ " _42 = _8 - _15\n",
+ " _43 = _41 * _42.T\n",
+ " _44 = _43.T @ _38\n",
+ " _45 = np.linalg.svd(_44, full_matrices=False)\n",
+ " _46 = np.array(0.0001) * _45[1][0]\n",
+ " _47 = _45[1] > _46\n",
+ " _48 = _47.astype(np.dtype(np.int32))\n",
+ " _49 = np.sum(_48)\n",
+ " _50 = _38 @ _45[2].T[:, :_49,]\n",
+ " _51 = _16 @ _50\n",
+ " return _51[:, :2,]\n",
"\n"
]
}
],
"source": [
+ "from egglog.exp.array_api_program_gen import *\n",
+ "import numpy as np\n",
"import inspect\n",
"\n",
+ "\n",
+ "egraph = EGraph()\n",
+ "fn_program = egraph.let(\n",
+ " \"fn_program\",\n",
+ " EvalProgram(ndarray_function_two_program(simplified_res, NDArray.var(\"X\"), NDArray.var(\"y\")), {\"np\": np}),\n",
+ ")\n",
+ "egraph.run(array_api_program_gen_schedule)\n",
+ "fn = egraph.extract(fn_program.as_py_object).value\n",
+ "\n",
"print(inspect.getsource(fn))"
]
},
@@ -1527,7 +1646,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "From there we can complete our work, by optimizing with numba and we can call with our original values:\n"
+ "From there we can complete our work, by optimizing with numba and we can call with our original values:"
]
},
{
@@ -1539,20 +1658,163 @@
"name": "stderr",
"output_type": "stream",
"text": [
- "/var/folders/xn/05ktz3056kqd9n8frgd6236h0000gn/T/egglog-9b40af4a-3b8a-4996-a78a-fd6284dbf541.py:56: NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (Array(float64, 2, 'C', False, aligned=True), Array(float64, 2, 'A', False, aligned=True))\n",
- " _45 = _33 @ _40[2].T[:, :_44]\n"
+ "/var/folders/xn/05ktz3056kqd9n8frgd6236h0000gn/T/egglog-39d7fedf-0738-45c9-8c02-1e5921b60da3.py:62: NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (Array(float64, 2, 'C', False, aligned=True), Array(float64, 2, 'A', False, aligned=True))\n",
+ " _50 = _38 @ _45[2].T[:, :_49,]\n"
]
},
{
"data": {
"text/plain": [
- "array([[ 0.64233002],\n",
- " [ 0.63661245],\n",
- " [-1.603293 ],\n",
- " ...,\n",
- " [-1.1506433 ],\n",
- " [ 0.71687176],\n",
- " [-1.51119579]])"
+ "array([[ 8.06179978e+00, 3.00420621e-01],\n",
+ " [ 7.12868772e+00, -7.86660426e-01],\n",
+ " [ 7.48982797e+00, -2.65384488e-01],\n",
+ " [ 6.81320057e+00, -6.70631068e-01],\n",
+ " [ 8.13230933e+00, 5.14462530e-01],\n",
+ " [ 7.70194674e+00, 1.46172097e+00],\n",
+ " [ 7.21261762e+00, 3.55836209e-01],\n",
+ " [ 7.60529355e+00, -1.16338380e-02],\n",
+ " [ 6.56055159e+00, -1.01516362e+00],\n",
+ " [ 7.34305989e+00, -9.47319209e-01],\n",
+ " [ 8.39738652e+00, 6.47363392e-01],\n",
+ " [ 7.21929685e+00, -1.09646389e-01],\n",
+ " [ 7.32679599e+00, -1.07298943e+00],\n",
+ " [ 7.57247066e+00, -8.05464137e-01],\n",
+ " [ 9.84984300e+00, 1.58593698e+00],\n",
+ " [ 9.15823890e+00, 2.73759647e+00],\n",
+ " [ 8.58243141e+00, 1.83448945e+00],\n",
+ " [ 7.78075375e+00, 5.84339407e-01],\n",
+ " [ 8.07835876e+00, 9.68580703e-01],\n",
+ " [ 8.02097451e+00, 1.14050366e+00],\n",
+ " [ 7.49680227e+00, -1.88377220e-01],\n",
+ " [ 7.58648117e+00, 1.20797032e+00],\n",
+ " [ 8.68104293e+00, 8.77590154e-01],\n",
+ " [ 6.25140358e+00, 4.39696367e-01],\n",
+ " [ 6.55893336e+00, -3.89222752e-01],\n",
+ " [ 6.77138315e+00, -9.70634453e-01],\n",
+ " [ 6.82308032e+00, 4.63011612e-01],\n",
+ " [ 7.92461638e+00, 2.09638715e-01],\n",
+ " [ 7.99129024e+00, 8.63787128e-02],\n",
+ " [ 6.82946447e+00, -5.44960851e-01],\n",
+ " [ 6.75895493e+00, -7.59002759e-01],\n",
+ " [ 7.37495254e+00, 5.65844592e-01],\n",
+ " [ 9.12634625e+00, 1.22443267e+00],\n",
+ " [ 9.46768199e+00, 1.82522635e+00],\n",
+ " [ 7.06201386e+00, -6.63400423e-01],\n",
+ " [ 7.95876243e+00, -1.64961722e-01],\n",
+ " [ 8.61367201e+00, 4.03253602e-01],\n",
+ " [ 8.33041759e+00, 2.28133530e-01],\n",
+ " [ 6.93412007e+00, -7.05519379e-01],\n",
+ " [ 7.68823131e+00, -9.22362309e-03],\n",
+ " [ 7.91793715e+00, 6.75121313e-01],\n",
+ " [ 5.66188065e+00, -1.93435524e+00],\n",
+ " [ 7.24101468e+00, -2.72615132e-01],\n",
+ " [ 6.41443556e+00, 1.24730131e+00],\n",
+ " [ 6.85944381e+00, 1.05165396e+00],\n",
+ " [ 6.76470393e+00, -5.05151855e-01],\n",
+ " [ 8.08189937e+00, 7.63392750e-01],\n",
+ " [ 7.18676904e+00, -3.60986823e-01],\n",
+ " [ 8.31444876e+00, 6.44953177e-01],\n",
+ " [ 7.67196741e+00, -1.34893840e-01],\n",
+ " [-1.45927545e+00, 2.85437643e-02],\n",
+ " [-1.79770574e+00, 4.84385502e-01],\n",
+ " [-2.41694888e+00, -9.27840307e-02],\n",
+ " [-2.26247349e+00, -1.58725251e+00],\n",
+ " [-2.54867836e+00, -4.72204898e-01],\n",
+ " [-2.42996725e+00, -9.66132066e-01],\n",
+ " [-2.44848456e+00, 7.95961954e-01],\n",
+ " [-2.22666513e-01, -1.58467318e+00],\n",
+ " [-1.75020123e+00, -8.21180130e-01],\n",
+ " [-1.95842242e+00, -3.51563753e-01],\n",
+ " [-1.19376031e+00, -2.63445570e+00],\n",
+ " [-1.85892567e+00, 3.19006544e-01],\n",
+ " [-1.15809388e+00, -2.64340991e+00],\n",
+ " [-2.66605725e+00, -6.42504540e-01],\n",
+ " [-3.78367218e-01, 8.66389312e-02],\n",
+ " [-1.20117255e+00, 8.44373592e-02],\n",
+ " [-2.76810246e+00, 3.21995363e-02],\n",
+ " [-7.76854039e-01, -1.65916185e+00],\n",
+ " [-3.49805433e+00, -1.68495616e+00],\n",
+ " [-1.09042788e+00, -1.62658350e+00],\n",
+ " [-3.71589615e+00, 1.04451442e+00],\n",
+ " [-9.97610366e-01, -4.90530602e-01],\n",
+ " [-3.83525931e+00, -1.40595806e+00],\n",
+ " [-2.25741249e+00, -1.42679423e+00],\n",
+ " [-1.25571326e+00, -5.46424197e-01],\n",
+ " [-1.43755762e+00, -1.34424979e-01],\n",
+ " [-2.45906137e+00, -9.35277280e-01],\n",
+ " [-3.51848495e+00, 1.60588866e-01],\n",
+ " [-2.58979871e+00, -1.74611728e-01],\n",
+ " [ 3.07487884e-01, -1.31887146e+00],\n",
+ " [-1.10669179e+00, -1.75225371e+00],\n",
+ " [-6.05524589e-01, -1.94298038e+00],\n",
+ " [-8.98703769e-01, -9.04940034e-01],\n",
+ " [-4.49846635e+00, -8.82749915e-01],\n",
+ " [-2.93397799e+00, 2.73791065e-02],\n",
+ " [-2.10360821e+00, 1.19156767e+00],\n",
+ " [-2.14258208e+00, 8.87797815e-02],\n",
+ " [-2.47945603e+00, -1.94073927e+00],\n",
+ " [-1.32552574e+00, -1.62869550e-01],\n",
+ " [-1.95557887e+00, -1.15434826e+00],\n",
+ " [-2.40157020e+00, -1.59458341e+00],\n",
+ " [-2.29248878e+00, -3.32860296e-01],\n",
+ " [-1.27227224e+00, -1.21458428e+00],\n",
+ " [-2.93176055e-01, -1.79871509e+00],\n",
+ " [-2.00598883e+00, -9.05418042e-01],\n",
+ " [-1.18166311e+00, -5.37570242e-01],\n",
+ " [-1.61615645e+00, -4.70103580e-01],\n",
+ " [-1.42158879e+00, -5.51244626e-01],\n",
+ " [ 4.75973788e-01, -7.99905482e-01],\n",
+ " [-1.54948259e+00, -5.93363582e-01],\n",
+ " [-7.83947399e+00, 2.13973345e+00],\n",
+ " [-5.50747997e+00, -3.58139892e-02],\n",
+ " [-6.29200850e+00, 4.67175777e-01],\n",
+ " [-5.60545633e+00, -3.40738058e-01],\n",
+ " [-6.85055995e+00, 8.29825394e-01],\n",
+ " [-7.41816784e+00, -1.73117995e-01],\n",
+ " [-4.67799541e+00, -4.99095015e-01],\n",
+ " [-6.31692685e+00, -9.68980756e-01],\n",
+ " [-6.32773684e+00, -1.38328993e+00],\n",
+ " [-6.85281335e+00, 2.71758963e+00],\n",
+ " [-4.44072512e+00, 1.34723692e+00],\n",
+ " [-5.45009572e+00, -2.07736942e-01],\n",
+ " [-5.66033713e+00, 8.32713617e-01],\n",
+ " [-5.95823722e+00, -9.40175447e-02],\n",
+ " [-6.75926282e+00, 1.60023206e+00],\n",
+ " [-5.80704331e+00, 2.01019882e+00],\n",
+ " [-5.06601233e+00, -2.62733839e-02],\n",
+ " [-6.60881882e+00, 1.75163587e+00],\n",
+ " [-9.17147486e+00, -7.48255067e-01],\n",
+ " [-4.76453569e+00, -2.15573720e+00],\n",
+ " [-6.27283915e+00, 1.64948141e+00],\n",
+ " [-5.36071189e+00, 6.46120732e-01],\n",
+ " [-7.58119982e+00, -9.80722934e-01],\n",
+ " [-4.37150279e+00, -1.21297458e-01],\n",
+ " [-5.72317531e+00, 1.29327553e+00],\n",
+ " [-5.27915920e+00, -4.24582377e-02],\n",
+ " [-4.08087208e+00, 1.85936572e-01],\n",
+ " [-4.07703640e+00, 5.23238483e-01],\n",
+ " [-6.51910397e+00, 2.96976389e-01],\n",
+ " [-4.58371942e+00, -8.56815813e-01],\n",
+ " [-6.22824009e+00, -7.12719638e-01],\n",
+ " [-5.22048773e+00, 1.46819509e+00],\n",
+ " [-6.80015000e+00, 5.80895175e-01],\n",
+ " [-3.81515972e+00, -9.42985932e-01],\n",
+ " [-5.10748966e+00, -2.13059000e+00],\n",
+ " [-6.79671631e+00, 8.63090395e-01],\n",
+ " [-6.52449599e+00, 2.44503527e+00],\n",
+ " [-4.99550279e+00, 1.87768525e-01],\n",
+ " [-3.93985300e+00, 6.14020389e-01],\n",
+ " [-5.20383090e+00, 1.14476808e+00],\n",
+ " [-6.65308685e+00, 1.80531976e+00],\n",
+ " [-5.10555946e+00, 1.99218201e+00],\n",
+ " [-5.50747997e+00, -3.58139892e-02],\n",
+ " [-6.79601924e+00, 1.46068695e+00],\n",
+ " [-6.84735943e+00, 2.42895067e+00],\n",
+ " [-5.64500346e+00, 1.67771734e+00],\n",
+ " [-5.17956460e+00, -3.63475041e-01],\n",
+ " [-4.96774090e+00, 8.21140550e-01],\n",
+ " [-5.88614539e+00, 2.34509051e+00],\n",
+ " [-4.68315426e+00, 3.32033811e-01]])"
]
},
"execution_count": 13,
@@ -1623,7 +1885,7 @@
"egraph = EGraph()\n",
"egraph.register(fn.compile())\n",
"egraph.run(program_gen_ruleset.saturate())\n",
- "print(egraph.eval(fn.statements))"
+ "print(egraph.extract(fn.statements).value)"
]
},
{
@@ -1642,7805 +1904,17 @@
"outputs": [
{
"data": {
- "image/svg+xml": [
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n"
- ],
- "text/html": [
- "\n",
- "\n",
- "\n",
- "\n",
- "\n",
- "\n"
- ],
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "2a45463df1ca40a2ac30f6f50da2b7b4",
+ "version_major": 2,
+ "version_minor": 1
+ },
"text/plain": [
- "EGraph(_flatted_deps=[], default_ruleset=None)"
+ "VisualizerWidget(egraphs=['{\"nodes\":{\"function-0-Program_statement\":{\"op\":\"·.statement\",\"children\":[\"function-…"
]
},
- "execution_count": 15,
"metadata": {},
- "output_type": "execute_result"
+ "output_type": "display_data"
}
],
"source": [
@@ -9582,7 +2056,7 @@
],
"metadata": {
"kernelspec": {
- "display_name": "Python 3 (ipykernel)",
+ "display_name": "egglog",
"language": "python",
"name": "python3"
},
@@ -9596,7 +2070,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.12.1"
+ "version": "3.13.3"
}
},
"nbformat": 4,
diff --git a/docs/explanation/indexing_pushdown.ipynb b/docs/explanation/indexing_pushdown.ipynb
index bea42460..508e7a8c 100644
--- a/docs/explanation/indexing_pushdown.ipynb
+++ b/docs/explanation/indexing_pushdown.ipynb
@@ -257,7 +257,7 @@
"\n",
"from egglog.exp.array_api import *\n",
"\n",
- "egraph = EGraph([array_api_module])\n",
+ "egraph = EGraph()\n",
"\n",
"\n",
"@egraph.register\n",
@@ -267,6 +267,7 @@
"\n",
"res = abs(NDArray.var(\"x\"))[NDArray.var(\"idx\")]\n",
"egraph.register(res)\n",
+ "egraph.run(array_api_schedule)\n",
"egraph.run(100)\n",
"egraph.display()\n",
"\n",
@@ -720,10 +721,10 @@
}
],
"source": [
- "egraph = EGraph([array_api_module])\n",
+ "egraph = EGraph()\n",
"\n",
"\n",
- "@egraph.function(cost=0)\n",
+ "@function(cost=0)\n",
"def value_abs(v: Value) -> Value:\n",
" \"\"\"Absolute value of a scalar value\"\"\"\n",
"\n",
@@ -738,7 +739,7 @@
" yield rewrite(abs(x).index(ti)).to(value_abs(x.index(ti)))\n",
"\n",
"\n",
- "@egraph.function(cost=0)\n",
+ "@function(cost=0)\n",
"def translate_index(x: NDArray, y: IndexKey, z: TupleInt) -> TupleInt:\n",
" \"\"\"Translates indexing `z` into `x[y]` into an indexing directly into `x`\"\"\"\n",
"\n",
@@ -752,12 +753,13 @@
" # Shape is more complicated and we will omit for now\n",
"\n",
"\n",
- "@egraph.function\n",
+ "@function\n",
"def an_index() -> TupleInt:\n",
" \"\"\"Some index into an array\"\"\"\n",
"\n",
"\n",
"egraph.register(res.shape, res.dtype, res.index(an_index()))\n",
+ "egraph.run(array_api_schedule)\n",
"egraph.run(100)\n",
"egraph.display()\n",
"\n",
diff --git a/docs/explanation/pldi_2023_presentation.ipynb b/docs/explanation/pldi_2023_presentation.ipynb
index 94acf518..508131db 100644
--- a/docs/explanation/pldi_2023_presentation.ipynb
+++ b/docs/explanation/pldi_2023_presentation.ipynb
@@ -429,7 +429,7 @@
"source": [
"## Open Source Data Science Ecosystem in Python\n",
"\n",
- "> The term “ecosystem” is often used to describe the modern open-source scientific software. In biology, the term “ecosystem” is defined as a biological community of interacting organisms and their physical environment. Modern open-source scientific software development occurs in a similarly interconnected and interoperable fashion.\n",
+ "> The term \u201cecosystem\u201d is often used to describe the modern open-source scientific software. In biology, the term \u201cecosystem\u201d is defined as a biological community of interacting organisms and their physical environment. Modern open-source scientific software development occurs in a similarly interconnected and interoperable fashion.\n",
"\n",
"from [Jupyter Meets the Earth: Ecosystem](https://jupytearth.org/jupyter-resources/introduction/ecosystem.html)\n",
"\n",
@@ -2996,7 +2996,6 @@
"egraph = EGraph()\n",
"\n",
"\n",
- "@egraph.class_\n",
"class Num(Expr):\n",
" @classmethod\n",
" def var(cls, name: StringLike) -> Num: ...\n",
@@ -4927,4 +4926,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
-}
+}
\ No newline at end of file
diff --git a/docs/reference/high-level.md b/docs/reference/high-level.md
index 5872ec77..8677b9bd 100644
--- a/docs/reference/high-level.md
+++ b/docs/reference/high-level.md
@@ -23,7 +23,6 @@ from egglog import *
egraph = EGraph()
-@egraph.class_
class Math(Expr):
def __init__(self, value: i64Like) -> None:
...
diff --git a/docs/reference/python-integration.md b/docs/reference/python-integration.md
index 466ebc0a..2789a51d 100644
--- a/docs/reference/python-integration.md
+++ b/docs/reference/python-integration.md
@@ -204,8 +204,8 @@ Alongside this, we support a function `dict_update` method, which can allow you
def my_add(a, b):
return a + b
-amended_globals = PyObject(globals()).dict_update("one", 1)
-evalled = py_eval("my_add(one, 2)", locals(), amended_globals)
+amended_globals = PyObject({"my_add": my_add}).dict_update("one", 1)
+evalled = py_eval("my_add(one, 2)", {}, amended_globals)
assert EGraph().extract(evalled).value == 3
```
diff --git a/docs/tutorials/getting-started.ipynb b/docs/tutorials/getting-started.ipynb
index 0644eae4..489c3dbe 100644
--- a/docs/tutorials/getting-started.ipynb
+++ b/docs/tutorials/getting-started.ipynb
@@ -83,7 +83,6 @@
"metadata": {},
"outputs": [],
"source": [
- "@egraph.class_\n",
"class Dim(Expr):\n",
" \"\"\"\n",
" A dimension of a matix.\n",
@@ -108,8 +107,8 @@
"tags": []
},
"source": [
- "As you can see, you must wrap any class with the `egraph.class_` to register\n",
- "it with the egraph and be able to use it like a Python class.\n",
+ "As you can see, you can define a class that inherits from `Expr` to create\n",
+ "a new type in the e-graph. The class will be automatically registered when you use it.\n",
"\n",
"### Testing in a notebook\n",
"\n",
@@ -117,7 +116,8 @@
"\n",
"```python\n",
"from matrix import *\n",
- "```\n"
+ "```\n",
+ ""
]
},
{
@@ -408,7 +408,8 @@
}
],
"source": [
- "egraph.simplify(res, 10)"
+ "egraph.run(10)\n",
+ "egraph.extract(res)"
]
},
{
@@ -429,7 +430,6 @@
"metadata": {},
"outputs": [],
"source": [
- "@egraph.class_\n",
"class Matrix(Expr):\n",
" @classmethod\n",
" def identity(cls, dim: Dim) -> Matrix:\n",
@@ -459,7 +459,7 @@
" \"\"\"\n",
"\n",
"\n",
- "@egraph.function\n",
+ "@function\n",
"def kron(a: Matrix, b: Matrix) -> Matrix:\n",
" \"\"\"\n",
" Kronecker product of two matrices.\n",
@@ -530,8 +530,9 @@
"x = Matrix.identity(Dim.named(\"x\"))\n",
"y = Matrix.identity(Dim.named(\"y\"))\n",
"x_mult_y = x @ y\n",
- "print(egraph.simplify(x_mult_y.ncols(), 10))\n",
- "print(egraph.simplify(x_mult_y.nrows(), 10))"
+ "egraph.run(10)\n",
+ "print(egraph.extract(x_mult_y.ncols()))\n",
+ "print(egraph.extract(x_mult_y.nrows()))"
]
},
{
@@ -1096,7 +1097,8 @@
")\n",
"# Create an example which should equal the kronecker product of A and B\n",
"ex1 = kron(Matrix.identity(n), B) @ kron(A, Matrix.identity(m))\n",
- "egraph.simplify(ex1, 20)"
+ "egraph.run(20)\n",
+ "egraph.extract(ex1)"
]
},
{
@@ -1209,7 +1211,8 @@
],
"source": [
"ex2 = kron(Matrix.identity(p), C) @ kron(A, Matrix.identity(m))\n",
- "egraph.simplify(ex2, 20)"
+ "egraph.run(20)\n",
+ "egraph.extract(ex2)"
]
},
{
@@ -1243,4 +1246,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
-}
+}
\ No newline at end of file
diff --git a/docs/tutorials/tut_2_datalog.py b/docs/tutorials/tut_2_datalog.py
index f1c479dd..66b2c942 100644
--- a/docs/tutorials/tut_2_datalog.py
+++ b/docs/tutorials/tut_2_datalog.py
@@ -150,7 +150,7 @@ def _(a: i64, b: i64, c: i64, ab: i64, bc: i64) -> Iterable[RewriteOrRule]:
# Let's run our rules and check we get the desired shortest path
egraph.run(run().saturate())
-egraph.check(path_len(1, 3) == 20)
+egraph.check(eq(path_len(1, 3)).to(20))
egraph
diff --git a/python/egglog/egraph.py b/python/egglog/egraph.py
index 819248d0..a0b1cbd5 100644
--- a/python/egglog/egraph.py
+++ b/python/egglog/egraph.py
@@ -1042,9 +1042,9 @@ def extract_multiple(self, expr: BASE_EXPR, n: int) -> list[BASE_EXPR]:
return [cast("BASE_EXPR", RuntimeExpr.__from_values__(self.__egg_decls__, expr)) for expr in new_exprs]
def _run_extract(self, expr: RuntimeExpr, n: int) -> bindings._CommandOutput:
- expr = self._state.typed_expr_to_egg(expr.__egg_typed_expr__)
+ egg_expr = self._state.typed_expr_to_egg(expr.__egg_typed_expr__)
# If we have defined any cost tables use the custom extraction
- args = (expr, bindings.Lit(span(2), bindings.Int(n)))
+ args = (egg_expr, bindings.Lit(span(2), bindings.Int(n)))
if self._state.cost_callables:
cmd: bindings._Command = bindings.UserDefined(span(2), "extract", list(args))
else:
@@ -1052,7 +1052,7 @@ def _run_extract(self, expr: RuntimeExpr, n: int) -> bindings._CommandOutput:
try:
return self._egraph.run_program(cmd)[0]
except BaseException as e:
- e.add_note("Extracting: " + str(expr))
+ e.add_note("while extracting expr:\n" + str(expr))
raise
def push(self) -> None:
diff --git a/src/egraph.rs b/src/egraph.rs
index 76949074..a8b253e2 100644
--- a/src/egraph.rs
+++ b/src/egraph.rs
@@ -64,20 +64,19 @@ impl EGraph {
cmds_str = cmds_str + &cmd.to_string() + "\n";
}
info!("Running commands:\n{}", cmds_str);
- let res = py.detach(|| {
- self.egraph.run_program(commands).map_err(|e| {
- WrappedError::Egglog(e, "\nWhen running commands:\n".to_string() + &cmds_str)
- })
- });
- if res.is_ok()
- && let Some(cmds) = &mut self.cmds
- {
- cmds.push_str(&cmds_str);
- }
+ let res = py.detach(|| self.egraph.run_program(commands));
if let Some(err) = PyErr::take(py) {
return Err(WrappedError::Py(err));
}
- res.map(|xs| xs.iter().map(|o| o.into()).collect())
+ match res {
+ Err(e) => Err(WrappedError::Egglog(e)),
+ Ok(outputs) => {
+ if let Some(cmds) = &mut self.cmds {
+ cmds.push_str(&cmds_str);
+ }
+ Ok(outputs.into_iter().map(|o| o.into()).collect())
+ }
+ }
}
/// Returns the text of the commands that have been run so far, if `record` was passed.
@@ -133,7 +132,7 @@ impl EGraph {
self.egraph
.eval_expr(&expr)
.map(|(s, v)| (s.name().to_string(), Value(v)))
- .map_err(|e| WrappedError::Egglog(e, format!("\nWhen evaluating expr: {expr}")))
+ .map_err(|e| WrappedError::Egglog(e))
});
if let Some(err) = PyErr::take(py) {
return Err(WrappedError::Py(err));
diff --git a/src/error.rs b/src/error.rs
index 8b468e11..5e225975 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -21,8 +21,7 @@ impl EggSmolError {
// https://pyo3.rs/latest/function/error_handling.html#foreign-rust-error-types
// TODO: Create classes for each of these errors
pub enum WrappedError {
- // Add additional context for egglog error
- Egglog(egglog::Error, String),
+ Egglog(egglog::Error),
ParseError(egglog::ast::ParseError),
Py(PyErr),
}
@@ -31,9 +30,7 @@ pub enum WrappedError {
impl From for PyErr {
fn from(error: WrappedError) -> Self {
match error {
- WrappedError::Egglog(error, str) => {
- PyErr::new::(error.to_string() + &str)
- }
+ WrappedError::Egglog(error) => PyErr::new::(error.to_string()),
WrappedError::Py(error) => error,
WrappedError::ParseError(error) => PyErr::new::(error.to_string()),
}
@@ -43,7 +40,7 @@ impl From for PyErr {
// Convert from an egglog::Error to a WrappedError
impl From for WrappedError {
fn from(other: egglog::Error) -> Self {
- Self::Egglog(other, String::new())
+ Self::Egglog(other)
}
}