torchzero 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

{torchzero-0.4.1.dist-info → torchzero-0.4.3.dist-info}/RECORD

@@ -1,54 +1,58 @@
-tests/test_identical.py,sha256=Y48_1f5WrltmO8a_-x-9Yltz2ZeMh8N8q3MGjOCkJhA,11552
+tests/test_identical.py,sha256=8Pw52Q19yeK5maYQEd2HYoOMItN599oRMUKzl-EugfQ,11550
 tests/test_module.py,sha256=qX3rjdSJsbA8JO17bPTUIDspe7bg2dogqxMw__KV7SU,2039
 tests/test_module_autograd.py,sha256=cncOlxtxmyJQHUd7nL9aWLRAr1kxtlKgVLqP3_qIb2E,21374
 tests/test_objective.py,sha256=HY0rK0z6PpiXvEsCu4mLgTlSVKusnT69S2GbuVcwMRo,7119
 tests/test_opts.py,sha256=hw7CCw7FD_RJSdiSacyXUSM7DI-_RfP8wJlsz079SNw,44263
 tests/test_tensorlist.py,sha256=B0Tq4_r-1DOYpS360X7IsLQiWn5fukhIMDKZM6zVO2Y,72164
 tests/test_utils_optimizer.py,sha256=_JoMqvXXZ6TxugS_CmfmP55Vvp0XrSPCjSz2nJJmaoI,8399
-torchzero/__init__.py,sha256=nit4KxrRoW6hJDGOy0jkphuawY5gAvPqrYY11Yct6fA,133
-torchzero/core/__init__.py,sha256=h9Ck7XX2XuJUTojU2IMa_2TprXZHbgo748txa3z7-2o,341
-torchzero/core/chain.py,sha256=dtFpxnw8vcbi3EeAANXyPtUmyPyv_VuZrTiPlLRmh7c,1899
-torchzero/core/functional.py,sha256=TSygtyQHDhqf998--hF48yIFr-y3Ycz8arjjR8x1ILU,3156
+torchzero/__init__.py,sha256=SZLJgf_sjHyqtTzz0f70AtHP_V_WloX1KQF8mm34zdg,175
+torchzero/_minimize/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchzero/_minimize/methods.py,sha256=1oktoSdWtiA0JEF34yTkY3_nPB5Q5ODHl18C0mcglNw,2445
+torchzero/_minimize/minimize.py,sha256=JJBmREQvhDxyqGM62xharsuebyefxRADkd6Bg_TE-DQ,17236
+torchzero/core/__init__.py,sha256=lufcll5r98gTjVfQSvz6-wfI0qMAgZtLLSByHuHTats,358
+torchzero/core/chain.py,sha256=-6vW-L5pzg2Rwpq3LKIAoqJGPvCkHKjt_B1boGikQmM,1900
+torchzero/core/functional.py,sha256=D125Hso8fHMSKlyhkir3GGJzXxuIitXmVhKn2Y9x-Ck,3272
 torchzero/core/modular.py,sha256=Xpp6jfiKArC3Q42G63I9qj3eWcYt-l7d-EIm-59ADcI,9584
-torchzero/core/module.py,sha256=HfbPfxXxgyBf9wQl7Fpw6B6Ux6UYfvPEmITC64ozb_Q,18012
+torchzero/core/module.py,sha256=DKGLwLWm9LkOBYZHW9QBoXo9eBgnYz7nmoCXJ0gl0e0,21210
 torchzero/core/objective.py,sha256=kEIlry7Bxf_zDUoqAIKUTRvvJmCEpn0Ad2crNt18GCc,40005
 torchzero/core/reformulation.py,sha256=UyAS_xq5sy_mMpmkvtwpHrZHd6Y2RgyPwN0zZlyxFTI,3857
-torchzero/core/transform.py,sha256=aJRBtvYjKqD-Ic_AkzeSINYDsTaBAErA-kocEl3PHZw,12244
+torchzero/core/transform.py,sha256=WlHoc5cCY1vXQrwMsIG0g3Kle93kBSbrBfxGz5X9_0Q,12251
 torchzero/linalg/__init__.py,sha256=wlry3dbncdsySKk6sSdiRefTcc8dIh4DcA0wFyU1MC8,407
 torchzero/linalg/benchmark.py,sha256=wiIMn-GY2xxWbHVf8CPbJddUPeUPq9OUDkvbp1iILYI,479
-torchzero/linalg/eigh.py,sha256=YC8x5NEWWsnc3suCebnTfeb4lVMhy-H8LGOZbGnwd8A,7902
+torchzero/linalg/eigh.py,sha256=l1fX_7hL-DFk8gu20-NuSKDJcRpz58KxUKQHeBhCcHE,9035
 torchzero/linalg/linalg_utils.py,sha256=1RBpQevb25_hbRieONw6CgnoWOJFXXv8oWTMugEivEk,385
-torchzero/linalg/linear_operator.py,sha256=mVEOvu6yY7TYhUdmZm1IAc6_pWnTaykKDgZu_-J-atk,16653
+torchzero/linalg/linear_operator.py,sha256=MWTY7DS8B8IkR28kVA9nmoM-OU-1eBsP22iYXkDrj9A,16654
 torchzero/linalg/matrix_power.py,sha256=gEWGvh3atc7745dwNcxNg0RtUrVgeKD6KxyRckKkkdQ,1255
-torchzero/linalg/orthogonalize.py,sha256=Fv6zv1JvS9AVwjiMVed55J8-pEbVZv7vqoEo5g0Zrv0,3270
+torchzero/linalg/orthogonalize.py,sha256=GSvDZA9evTpu3obqCkEocgpDp_91sRexoAwH2q0zTEY,5345
 torchzero/linalg/qr.py,sha256=KykXhSlye0vhyP5JjX6pkPnheHKLLbAKmDff8Hogxyo,2857
+torchzero/linalg/sketch.py,sha256=dKD9t7I7stv089cCvZyPAOZ0D9wzVG1TmV3297w0tk4,1261
 torchzero/linalg/solve.py,sha256=kING1WCioof8_EKgHeyr53dlft_9KtlJnwOWega3DnA,14355
 torchzero/linalg/svd.py,sha256=jmunSxM-twR5VCUGI_gmV3j7QxMJIe1aBoBlJf5i2fo,1432
 torchzero/linalg/torch_linalg.py,sha256=brhMXUZyYuxuEV-FyQerep1iL7auutW5kmgJpOzUROw,6001
-torchzero/modules/__init__.py,sha256=dsOalCw-OVkD8rhpQdcODc3Hsd_sQ2_2xVC-J8mlSuk,632
+torchzero/modules/__init__.py,sha256=ZN20E2ES6zDf5DuFbZpuCKFinFc5eGR1h00iYZ_XBGU,652
 torchzero/modules/opt_utils.py,sha256=aj7xqHmeze4izxG9k3L6ziG-K_yj8n8fkFpIv-X8V78,8141
 torchzero/modules/adaptive/__init__.py,sha256=X8w2Dal3k0WpLQN-WolnWBBgUyIiZF5RnqBlN0dcAYw,1081
-torchzero/modules/adaptive/adagrad.py,sha256=hMT-Al-vtD6tzPUpQ79LCNko97D7rJN5ji9JOfBqR3k,12015
+torchzero/modules/adaptive/adagrad.py,sha256=NDwmUZaEk0lWnbgYxN23yTWK5A5dQ9BtoKzRTFSKozY,12131
 torchzero/modules/adaptive/adahessian.py,sha256=ucf8loS_lU9VjCb_M42WwXESjPJ_KFChLGkIMFWXO5o,8734
-torchzero/modules/adaptive/adam.py,sha256=Okm7Sc9fMArQAZ7Ph4Etq68uL-IXKY4YNqHWpTzPoTY,3767
-torchzero/modules/adaptive/adan.py,sha256=965tBUwKy6uDiY2la6fVcGcsvGMs90Zg-ZHPtozJGe4,4110
+torchzero/modules/adaptive/adam.py,sha256=RDHYyIAJdi1Pxny8HOHiCFgvPztNwlJlCtzE_ZE-138,3896
+torchzero/modules/adaptive/adan.py,sha256=tmQHiJ5MNwOGP3fp479goHh0xXlhnzULhHxKcVZOkvM,4219
 torchzero/modules/adaptive/adaptive_heavyball.py,sha256=iDiZqke6z6FOR9mhoHMLMm7jvxjzHIQANTe0FBwNj1Q,2230
 torchzero/modules/adaptive/aegd.py,sha256=WLN6vvbSRhQ1P753M3bx_becSF-3cTbu37nhz3NvdGM,1903
 torchzero/modules/adaptive/esgd.py,sha256=gnah-7zk_fMsn7yIWivqDgnaaSdDFXpxg33ywF6TMZg,6173
-torchzero/modules/adaptive/ggt.py,sha256=eYCeV3GArdLv9WuWeim0V3CHJYl3FVKtrtsGshkqwWg,6608
-torchzero/modules/adaptive/lion.py,sha256=H3aI2qnrMtmkvXcoddzjjxdkoD5cq_QwIkLmd_bVPso,1085
+torchzero/modules/adaptive/ggt.py,sha256=7G0Hh8lWy4o73VYVHcZ1JJyDqeKcXi2Y6Qp3qIHosOY,6512
+torchzero/modules/adaptive/lion.py,sha256=yeKUt3WIITtWx97IQzudgbdai77MCfnL_cu90vRkTmA,1141
 torchzero/modules/adaptive/lre_optimizers.py,sha256=AwWUIwnBrozR2HFYLfJnMCBHAWWMKzkS63xFKstRgc0,9760
-torchzero/modules/adaptive/mars.py,sha256=w-cK-1tFuR74SY01xS5jsg1b9qs3l8eOptGrUyQ2m80,2261
+torchzero/modules/adaptive/mars.py,sha256=WquKzTnCZcxzslcvSBMFJVz_kjuCuAzlesw1bHnKqOg,2325
 torchzero/modules/adaptive/matrix_momentum.py,sha256=YefF2k746ke7qiiabdhCPCUFB1_fRddAfGCyIOwV3Ok,6789
-torchzero/modules/adaptive/msam.py,sha256=nqwjuhBMX2UO-omUIeOcD5ti6PIKfKs-RVCn7ourkKA,6946
+torchzero/modules/adaptive/msam.py,sha256=cHfdNkk3Joy2aENwUZXGf3N0P7zcxYGKuySf699OTfM,7051
 torchzero/modules/adaptive/muon.py,sha256=jQ6jlfM4vVRidGJ7FrLtgPnZeuIfW_zU72o7LvOKqh8,8023
 torchzero/modules/adaptive/natural_gradient.py,sha256=8UzacvvIMbYVVE2q0HQ9DLLHYlm1eu6cAiRsOv5XRzQ,7078
 torchzero/modules/adaptive/orthograd.py,sha256=0u2sfGZJjlJItLX2WRP5fLAD8Wd9SgJzQYAUpARJ64A,1813
-torchzero/modules/adaptive/rmsprop.py,sha256=qWVkRmUQ3dui9yBVYtAEll7OlXZDKNT_m70FakTOrTY,4529
-torchzero/modules/adaptive/rprop.py,sha256=a4_UkWse5u2JFAEIlxQqDBUwvUfxh1kNs2ZIhtccnWE,11540
+torchzero/modules/adaptive/rmsprop.py,sha256=sb709Smpkm8H3vYOsh7BzWni5hAf3nBQevhagyOt7mo,4655
+torchzero/modules/adaptive/rprop.py,sha256=vw-Rufa-gpHgq1gDarmNQexrFr13lPLq_mj3c3pNB_Q,11593
 torchzero/modules/adaptive/sam.py,sha256=CTMCqaH9s5EmKQyj1GpqSeTO1weyfsNWPYFN1xaSm_o,5709
-torchzero/modules/adaptive/shampoo.py,sha256=C_Mo7UFQtDxW4McWJjT731FNAp3g9MqF0Hka54Yi3xQ,9847
-torchzero/modules/adaptive/soap.py,sha256=hz2N6-jUSWU93RNViIS1c-Ue2uKmQx6BxyYg6mEa2fo,12408
+torchzero/modules/adaptive/shampoo.py,sha256=1WpjroFS37HmDLV51iK4d8vtnJWFrGCsDkoQav0p47E,10048
+torchzero/modules/adaptive/soap.py,sha256=jyS6F2o4bMKzMU8H2dDggFQEqMqw4W1rX78u8p3uaV4,12619
 torchzero/modules/adaptive/sophia_h.py,sha256=O_izgGlUgUlpH3Oi5PdCKTyxus4yO1PaJUFhGXuGG9k,7063
 torchzero/modules/adaptive/psgd/__init__.py,sha256=g73mAkWEutwU6jzjiwdbYk5Yxgs4i6QVWefFKkm8cDw,223
 torchzero/modules/adaptive/psgd/_psgd_utils.py,sha256=YtwbUKyVWITZPmpwCBJBC42XQP9HcxNx_znEaIv3hsI,1096
@@ -58,21 +62,20 @@ torchzero/modules/adaptive/psgd/psgd_kron_newton.py,sha256=oH-oI1pvbR-z6H6ma1O2G
 torchzero/modules/adaptive/psgd/psgd_kron_whiten.py,sha256=vmhkY6cKaRE5qzy_4tUkIJp6qC3L6ESZMuiU_ih5tR4,7299
 torchzero/modules/adaptive/psgd/psgd_lra_newton.py,sha256=JL8JmqHgcFqfkX7VeD3sRvNj0xeCuDTHxjNyQ_HigBw,4709
 torchzero/modules/adaptive/psgd/psgd_lra_whiten.py,sha256=SaNYtE4_2tV29CbVaTHi8A6RxmhoMaucF5NoMRg6QaA,4197
+torchzero/modules/basis/__init__.py,sha256=MeXoykwqqmWt-Gx8YWMycVL7m5N4j7Ob_L0GbcwLOfM,65
+torchzero/modules/basis/ggt_basis.py,sha256=NVddZrv58lm7M2Q2j5_3YYLcBYRdeSB_y03bxExSiJs,7772
+torchzero/modules/basis/soap_basis.py,sha256=pwlxIa9lW9V1NcLPmhm--LVbyq7ALSfkV_4b6ki1hO8,10479
 torchzero/modules/clipping/__init__.py,sha256=ZaffMF7mIRK6hZSfuZadgjNTX6hF5ANiLBny2w3S7I8,250
 torchzero/modules/clipping/clipping.py,sha256=C2dMt0rpuiLMsKq2EWi8qhISSxfCU0nKKGgjWEk2Yxc,14198
-torchzero/modules/clipping/ema_clipping.py,sha256=D4NgXzXYMjK_SKQU3rVoOKzaCd9igGQg_7sXiGMgMqI,6750
-torchzero/modules/clipping/growth_clipping.py,sha256=I1nk5xXBjk0BzWYzMC58LZHouY44myZNIUjM-duv7zc,6508
+torchzero/modules/clipping/ema_clipping.py,sha256=7lFkQWVkchxlZynYXS4JDjhxB8T5tbE0qsP3GXK6mrA,6916
+torchzero/modules/clipping/growth_clipping.py,sha256=VAmUUeIsSGWrGmZiFAngWUBBsxj4d0QAMf36oAMZL8A,6556
 torchzero/modules/conjugate_gradient/__init__.py,sha256=G5WcVoiQYupRBeqjI4lCraGeXNSvWT-_-ynpcE6NQS8,184
 torchzero/modules/conjugate_gradient/cg.py,sha256=fcmP77_v_RPpb0sDV2B_90FvFY67FdJt54KHdccY5YU,14540
-torchzero/modules/experimental/__init__.py,sha256=YbBrWu2vkXHiBcDXmim-Yte4ZxfmQCs_0fCeIArvtnM,942
-torchzero/modules/experimental/adanystrom.py,sha256=fUWPxxi1aJhWme_d31dBG0XxEZY1hJr6AEiFHdFDxCQ,8970
-torchzero/modules/experimental/common_directions_whiten.py,sha256=R_1fQKlvMD99oFrflJLgxl6ObV8jyPc7-NxAUFQeoYA,4941
-torchzero/modules/experimental/coordinate_momentum.py,sha256=HzKy8X5qEvud-xKHJYHpzH6ObxzvYcMcdgodsCw4Bbk,1099
-torchzero/modules/experimental/cubic_adam.py,sha256=RhcHajUfUAcXZDks0X0doR18YtMItQYPmxuEihud4bo,5137
+torchzero/modules/experimental/__init__.py,sha256=najUDh01Av6gEeMYRV9X9lWAr4ZrC6ZgJcPtNpon7ZQ,734
+torchzero/modules/experimental/coordinate_momentum.py,sha256=4BMmgooPysYlX7QOaTUjBn6MNfBAMujM5TCm72vSexw,1152
+torchzero/modules/experimental/cubic_adam.py,sha256=97sgbtkqG1ziXOMxlCor-L-UzzqgSumz8shVOgYL4oQ,5303
 torchzero/modules/experimental/curveball.py,sha256=beHGD1Wh9GxYqMBh1k9Ru6TG3U9eZR6_l8ZUQcZzYxw,2765
 torchzero/modules/experimental/dct.py,sha256=CW-Y2gcjlHlxtIx7SekUOfw2EzujA6v0LcjDYGAfh6M,2433
-torchzero/modules/experimental/eigen_sr1.py,sha256=rCcWVplTWQh91xpgDap35CGEex41C19irUfDlq9lviU,6865
-torchzero/modules/experimental/eigengrad.py,sha256=UPuyo-OmCmu3XLAPclIfsnMN4qcNwX83m7S_55syukA,8455
 torchzero/modules/experimental/fft.py,sha256=s95EzvK4-ZJdwZbVhtqwirY9eVy7v6mFDRMgoLY9wjo,3020
 torchzero/modules/experimental/gradmin.py,sha256=LajM0GU1fB6PsGDg8k0KjKI73RvyZYqPvzcdoVYDq-c,3752
 torchzero/modules/experimental/higher_order_newton.py,sha256=qLSCbkmd7dw0lAhOJGpvvOesZfCMNt2Vz_mc7HknCMQ,12131
@@ -82,15 +85,15 @@ torchzero/modules/experimental/newton_solver.py,sha256=aHZh8EA-QQop3iGz7Ge37KTNg
 torchzero/modules/experimental/newtonnewton.py,sha256=TYUuQwHu8bom08czU9lP7MQq5qFBq_JYZTH_Wmm4g-o,3269
 torchzero/modules/experimental/reduce_outward_lr.py,sha256=ehctg5zLEOHPfiQQUq5ShMj3pDhtxqdNUEneMR9l7Bs,1275
 torchzero/modules/experimental/scipy_newton_cg.py,sha256=psllNtDwUbkVAXBDKwWEueatOmDNPFy-pMwBkqF3_r0,3902
-torchzero/modules/experimental/spsa1.py,sha256=DiQ_nHAC8gnqoNNK7oe6djOiwpwvI5aPtpKA43F7jrQ,3607
 torchzero/modules/experimental/structural_projections.py,sha256=IwpgibNDO0slzMyi6djQXRhQO6IagUgUUCr_-7US1IE,4104
-torchzero/modules/grad_approximation/__init__.py,sha256=_mQ2sWvnMfqc3RQcVmZuBlphtLZCO7z819abGY6kYuM,196
+torchzero/modules/grad_approximation/__init__.py,sha256=BAFXc73_ORySVDyXiyZxpusXWn7K66KFT9LZEMwVKes,221
 torchzero/modules/grad_approximation/fdm.py,sha256=hq7U8UkzCfc7z0J1ZmZo9xOLzHHY0uRjebcwZQrBCzA,4376
 torchzero/modules/grad_approximation/forward_gradient.py,sha256=7fKZoKetYzgD85L3W0x1oG56SdWHj5MDWwmWpV7bpr4,3949
 torchzero/modules/grad_approximation/grad_approximator.py,sha256=hX4nqa0yw1OkA2UKmzZ3HhvMfL0Wwv1yQePxrgAueS8,4782
 torchzero/modules/grad_approximation/rfdm.py,sha256=-5zqMB98YNNa1aQXXtf6UNGSJxySO7mn1NksWyPzp3o,19607
+torchzero/modules/grad_approximation/spsa1.py,sha256=DiQ_nHAC8gnqoNNK7oe6djOiwpwvI5aPtpKA43F7jrQ,3607
 torchzero/modules/least_squares/__init__.py,sha256=mJwE2IXVB3mn_7BzsmDNKhfyViCV8GOrqHJJjz04HR4,41
-torchzero/modules/least_squares/gn.py,sha256=3RQ_7e35Ql9uVUUPi34nef9eQNeZ09fldi964V61Tgg,7889
+torchzero/modules/least_squares/gn.py,sha256=hufsWNq_UdEPFDFKNGgCiM4R9739Xu8JqYWSwKkdSZ8,8087
 torchzero/modules/line_search/__init__.py,sha256=_QjxUJmNC8OqtUuyTJp9wDfHNFKZBZqj6lttWKhG-cI,217
 torchzero/modules/line_search/_polyinterp.py,sha256=i3sNl6SFAUJi4oxhhjBlcxJY9KRunIZjJ8sGdaJOVjc,10990
 torchzero/modules/line_search/adaptive.py,sha256=YNabP6-01dhAUDAOuHRPZCwiV5xTRdHmkN667HQ6V3w,3798
@@ -102,21 +105,21 @@ torchzero/modules/line_search/strong_wolfe.py,sha256=9jGjxebuXHbl8wEFpvV0s4mMX4J
 torchzero/modules/misc/__init__.py,sha256=UYY9CeNepnC8H1LnFa829ux5MEjtGZ9zql624IbCFX8,825
 torchzero/modules/misc/debug.py,sha256=wFt9wB6IdRSsOGLhQjdjmGt4KdB0V5IT0iBFMj97R3Y,1617
 torchzero/modules/misc/escape.py,sha256=c_OMf2jQ7MbxkrXWNmgIpZrBe28N9f89tnzuCQ3fu3A,1930
-torchzero/modules/misc/gradient_accumulation.py,sha256=Xzjt_ulm6Z3mpmtagoUqoefhoeSDVnmX__tVbcI_RQE,2271
+torchzero/modules/misc/gradient_accumulation.py,sha256=1BVqGXwv1YPg7DRJWP0XY6s-vzxrvyXLdruM1Y5KJ5s,2326
 torchzero/modules/misc/homotopy.py,sha256=oa0YFYfv8kkg9v7nukdjTwinuyQa4Nt7kTpddUVCSKg,2257
-torchzero/modules/misc/misc.py,sha256=f-3qxBq1KYI3iGYJXzv1cHEJHc0ScEp-vCLCgiaEgJQ,15002
+torchzero/modules/misc/misc.py,sha256=eWVyYSYiQxcS7G7aVM4nFYiF0csE9qcztTaP4id5CbE,15306
 torchzero/modules/misc/multistep.py,sha256=twdE-lU9Wa0b_uquH9kZ-1OwP0gqWfFMJkdjVWJRwe4,6599
 torchzero/modules/misc/regularization.py,sha256=MCd_tnBYfFnx0b3sM1vHNQ_WbTVfo7l8pxmxGVgWcc0,5935
 torchzero/modules/misc/split.py,sha256=rmi9PgMgiqddrr8fY8Dbdcl2dgwTn9YBAve_bg5Zd08,4288
 torchzero/modules/misc/switch.py,sha256=_ycuD23gR0ZvIUmX3feYBr0_WTX22Pfhu3whpiSCMv4,3678
 torchzero/modules/momentum/__init__.py,sha256=AKWC4HIkN9ZJwN38dJvVJkFEhiP9r93G-kMDokBfsj8,281
-torchzero/modules/momentum/averaging.py,sha256=Q6WLwCJwgNY96YIfQXWpsX-2kDR7n0IOMDfZMvNVc9U,3035
+torchzero/modules/momentum/averaging.py,sha256=OTO_LRNiAhbcKTXrWI-uENqIOH_3DX5_1uYJ3eMVcJY,3202
 torchzero/modules/momentum/cautious.py,sha256=1hD2H08OQaNZG52sheRADBsuf9uJsaoLV4n-UVGUH3Y,8379
-torchzero/modules/momentum/momentum.py,sha256=MPHd4TU1bSlEKLGfueNdmaZ13V5J1suW6agBc3SvrTs,4389
-torchzero/modules/ops/__init__.py,sha256=xUYzWWLlSwaT8sw3dWywkALqI6YGCZgptWQJVy83HhM,1249
-torchzero/modules/ops/accumulate.py,sha256=f-Uutg7gNFRobTc5YI9JlfFiSacXmg0gDhIwQNwZSZg,3439
+torchzero/modules/momentum/momentum.py,sha256=aJ8o3gB9HebM9kutpadC5wI0MgMjn-c3J4GF3Z_n0Oc,4484
+torchzero/modules/ops/__init__.py,sha256=p5hwECuODOv6E4H0lETQHweSsUtMlsGE0d8bfTv2Rwc,1225
+torchzero/modules/ops/accumulate.py,sha256=mbJFwykU2fa6IIfsHVXdhmRp7QX1czpCWjw6AYkNn1k,3636
 torchzero/modules/ops/binary.py,sha256=eB6zwz5ZSSyeWvwVfuOFMjem93oMB7hCo4kNF705jn8,12219
-torchzero/modules/ops/higher_level.py,sha256=cUh-908S0GWVGekmUN5c_Vx0HP3P2tQoKN3COQM5TaQ,8965
+torchzero/modules/ops/higher_level.py,sha256=f9DFNI9rnxc-rShAJOfsiwvyGsWu8FsJwJf5yg_V4eg,9366
 torchzero/modules/ops/multi.py,sha256=WzNK07_wL7z0Gb2pmv5a15Oss6tW9IG79x1c4ZPmOqQ,8643
 torchzero/modules/ops/reduce.py,sha256=SzpkNV5NTsVFp-61a1m8lDKJ1ivJmfQofolFWxbbAe4,6526
 torchzero/modules/ops/unary.py,sha256=vXvWfDFo2CBFwb1ej_WV-fGg61lQRbwN4HklAik8tJY,4844
@@ -136,12 +139,12 @@ torchzero/modules/restarts/__init__.py,sha256=7282ePwN_I0vSeLPYS4TTclE9ZU7pL6Upy
 torchzero/modules/restarts/restars.py,sha256=gcRZ8VHGg60cFVzsk0TWa6-EXoqEFbEeP1p7fs2Av0Q,9348
 torchzero/modules/second_order/__init__.py,sha256=42HeVA3Azl_tXV0_injU-q4QOu7lXzt6AVUcwnPy4Ag,313
 torchzero/modules/second_order/ifn.py,sha256=oAjfFVjLzG6L4n_ELXAWGZSicWizilQy_hQf4hmOoL0,2019
-torchzero/modules/second_order/inm.py,sha256=OddoZHQfSuFnlx_7Zj2qiVcC2A_9yMVn_0Gy1A7hNAg,3420
+torchzero/modules/second_order/inm.py,sha256=_FnaUHKLl46AtI_XYwF52wtOUbAaO5EMUNRJspX5FEM,3574
 torchzero/modules/second_order/multipoint.py,sha256=mHG1SFLsILELIspxZ8U_hxJBlkGwzvUWg96bOIrQsIY,7500
-torchzero/modules/second_order/newton.py,sha256=QcLXsglvf4zJEwR4cldsGVZCABQtxb6U5qVmU3spN_A,11061
-torchzero/modules/second_order/newton_cg.py,sha256=k8G8CSmeIQZObkWVURFnbF_4g2UvJiwh3xToxn7sFJE,14816
-torchzero/modules/second_order/nystrom.py,sha256=WQFfJj0DOfWXyyx36C54m0WqZPIvTTK7n8U7khLhGLg,13359
-torchzero/modules/second_order/rsn.py,sha256=9s-JyJNNeDlIFv8YVGn7y8DGPnP93WJEjpUQXehX3uY,9980
+torchzero/modules/second_order/newton.py,sha256=W37_ePdAB1wnlRrNRd2ovNgkbodK1JV8J4SJytVuF_M,11456
+torchzero/modules/second_order/newton_cg.py,sha256=gHmpLRQ2FRr0750gYkFQ7XweJVZmYI6yG9H2vrKvAdA,14925
+torchzero/modules/second_order/nystrom.py,sha256=lGLjtzq2WAWcaT3E6Say82ySZ1yp9I2ASuOqyNTUmiQ,13361
+torchzero/modules/second_order/rsn.py,sha256=13t42cUvY8JQMC4zf4UsqKvpnTXuXZUZJDECCxRYWjg,11286
 torchzero/modules/smoothing/__init__.py,sha256=RYxCLLfG2onBbMUToaoedsr20rXaayyBt7Ov8OxULrU,80
 torchzero/modules/smoothing/laplacian.py,sha256=1cewdvnneKn51bbIBqKij0bkveKE7wOYCZ-aGlqzK5M,5201
 torchzero/modules/smoothing/sampling.py,sha256=bCH7wlTYZ_vtKUKSkI6znORxQ5Z6DGcpo10F-GYvFlE,12880
@@ -155,7 +158,7 @@ torchzero/modules/trust_region/cubic_regularization.py,sha256=QJjLRkfERvOzV5dTdy
 torchzero/modules/trust_region/dogleg.py,sha256=zwFR49gghxztVGEETF2D4AkeGgHkQRbHGGelav3GuFg,3619
 torchzero/modules/trust_region/levenberg_marquardt.py,sha256=-qbeEW3qRKou48bBdZ-u4Nv43TMt475XV6P_aWfxtqE,5039
 torchzero/modules/trust_region/trust_cg.py,sha256=X9rCJQWvptjZVH2H16iekvAYmleKQAYZKRKC3V0JjFY,4455
-torchzero/modules/trust_region/trust_region.py,sha256=oXMNIvboz0R_1J0Gfd4IvbnwZFl32csNVv-lTYGB0zk,12913
+torchzero/modules/trust_region/trust_region.py,sha256=ax1pJDr3NPLfojUXRMb-hsxD4MpQL1bPAOwozAVTCJI,12930
 torchzero/modules/variance_reduction/__init__.py,sha256=3pwPWZpjgz1btfLJ3rEaK7Wl8B1pDh0HIf0kvD_NJH8,22
 torchzero/modules/variance_reduction/svrg.py,sha256=hXEJ0PUYSksHV0ws3t3cE_4MUTTEn1Htu37iZdDdJCs,8746
 torchzero/modules/weight_decay/__init__.py,sha256=zQrjSujD0c-rKfKjUpuutfAODljsz1hS3zUNJW7zbh4,132
@@ -195,15 +198,15 @@ torchzero/utils/derivatives.py,sha256=Sc20EH2v2czjH9Z8UChvq0EaYtvOEJKEYOk3fVb0Z6
 torchzero/utils/metrics.py,sha256=XPpOvY257tb4mN3Sje1AVNlQkOXiW24_lXXdtd0JYok,3130
 torchzero/utils/numberlist.py,sha256=iMoqz4IzXy-aE9bqVYJ21GV6pl0z-NeTsXR-LaI8C24,6229
 torchzero/utils/optimizer.py,sha256=G741IvE57RaVYowr9FEqfRm_opPAeu4UWKU5iPKDMFA,8415
-torchzero/utils/optuna_tools.py,sha256=F-1Xg0n_29MVEb6lqgUFFNIl9BNJ6MOdIJPduoNH4JU,1325
-torchzero/utils/params.py,sha256=nQo270aOURU7rJ_D102y2pSXbzhJPK0Z_ehx4mZBMes,5784
+torchzero/utils/optuna_tools.py,sha256=t64nwyuIVP7xgeGVvIGMFBij2j5clhjY4BHtGEnyPVI,1323
+torchzero/utils/params.py,sha256=-amJs518rpI0zzYavTlWrl60JNrgsk1xxdGvIrSw1ZI,6406
 torchzero/utils/python_tools.py,sha256=HATghTNijlQxmw8rzJfZPPGj1CjcnRxEwogmrgqnARU,4577
-torchzero/utils/tensorlist.py,sha256=4rN8gm967pPmtO5kotXqIX7Mal0ps-IHkGBybfeWY4M,56357
+torchzero/utils/tensorlist.py,sha256=wpzBJvIAmw9VXsg1UF8gZtq-eh7GlvdM6WL_7NyPYlY,56363
 torchzero/utils/thoad_tools.py,sha256=G8k-z0vireEUtI3A_YAR6dtwYjSnN49e_GadcHwwQKc,2319
 torchzero/utils/torch_tools.py,sha256=DsHaSRGZ3-IuySZJTrkojTbaMMlttJFe0hFvB2xnl2U,5069
 torchzero/utils/benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-torchzero/utils/benchmarks/logistic.py,sha256=RHsjHEWkPqaag0kt3wfmdddh4DhftcyW9r70tj9OGp4,4382
-torchzero-0.4.1.dist-info/METADATA,sha256=hB0rFqXnaRbwVkFRwTwjXpKnIFLi8MBvLXbgXTuUGWk,564
-torchzero-0.4.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-torchzero-0.4.1.dist-info/top_level.txt,sha256=ETW_iE2ubg0oMyef_h-ayB5i1OOZZd4SNdR3ltIbHe0,16
-torchzero-0.4.1.dist-info/RECORD,,
+torchzero/utils/benchmarks/logistic.py,sha256=1c9kB6tDaKsSNlQn44_Lso2_g-85fQK45RvwLZOcJOo,4587
+torchzero-0.4.3.dist-info/METADATA,sha256=39RK0MpaBQIm0GpIK2YRwoeY5zegEBnJHCZIY4ExQ5k,564
+torchzero-0.4.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+torchzero-0.4.3.dist-info/top_level.txt,sha256=ETW_iE2ubg0oMyef_h-ayB5i1OOZZd4SNdR3ltIbHe0,16
+torchzero-0.4.3.dist-info/RECORD,,

torchzero/modules/experimental/adanystrom.py

@@ -1,258 +0,0 @@
-# pylint: disable = non-ascii-name
-import torch
-
-from ...core import Chainable, TensorTransform
-from ...linalg import (
-    OrthogonalizeMethod,
-    orthogonalize,
-    regularize_eigh,
-    torch_linalg,
-)
-from ...linalg.linear_operator import Eigendecomposition
-from ..adaptive.lre_optimizers import LREOptimizerBase
-from .eigengrad import _eigengrad_update_state_, eigengrad_apply
-
-
-def weighted_eigen_plus_rank1_mm(
-    # A1 = Q1 @ diag(L1) @ Q1.T
-    L1: torch.Tensor,
-    Q1: torch.Tensor,
-
-    # K2 = v2 @ v2.T
-    v2: torch.Tensor,
-
-    # second matrix
-    B: torch.Tensor,
-
-    # weights
-    w1: float,
-    w2: float,
-
-) -> torch.Tensor:
-    """
-    Computes ``(w1 * A1 + w2 * A2) @ B``, where ``A1`` is an eigendecomposition, ``A2`` is symmetric rank 1.
-
-    Returns ``(n, k)``
-
-    Args:
-        L1 (torch.Tensor): eigenvalues of A1, shape ``(rank,)``.
-        Q1 (torch.Tensor): eigenvectors of A1, shape ``(n, rank)``.
-        v2 (torch.Tensor): vector such that ``v v^T = A2``, shape ``(n,)``.
-        B (torch.Tensor): shape ``(n, k)``.
-        w1 (float): weight for A1.
-        w2 (float): weight for A2.
-
-    """
-    # sketch A1
-    QTB = Q1.T @ B # (rank, k)
-    LQTB = L1.unsqueeze(1) * QTB  # (rank, k)
-    sketch1 = Q1 @ LQTB  # (n, k)
-
-    # skecth A2
-    vB = v2 @ B
-    sketch2 = v2.outer(vB)
-
-    return w1 * sketch1 + w2 * sketch2
-
-
-def adanystrom_update(
-    L1: torch.Tensor,
-    Q1: torch.Tensor,
-    v2: torch.Tensor,
-    w1: float,
-    w2: float,
-    oversampling_p: int,
-    rank: int,
-    eig_tol: float,
-    damping: float,
-    rdamping: float,
-    orthogonalize_method: OrthogonalizeMethod,
-
-) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-    """computes the Nyström approximation of ``(w1 * A1 + w2 * A2)``,
-    where ``A1`` is an eigendecomposition, ``A2`` is symmetric rank 1.
-
-    returns L of shape ``(k, )`` and Q of shape ``(n, k)``.
-
-    Args:
-        L1 (torch.Tensor): eigenvalues of A1, shape ``(rank,)``.
-        Q1 (torch.Tensor): eigenvectors of A1, shape ``(n, rank)``.
-        v2 (torch.Tensor): vector such that ``v v^T = A2``, shape ``(n,)`` or ``(n, 1)``.
-        w1 (float): weight for A1.
-        w2 (float): weight for A2.
-    """
-    n = Q1.shape[0]
-    device = Q1.device
-    dtype = Q1.dtype
-    l = rank + oversampling_p
-
-    # gaussian test matrix
-    Omega = torch.randn(n, l, device=device, dtype=dtype)
-
-    # sketch
-    AOmega = weighted_eigen_plus_rank1_mm(L1, Q1, v2, Omega, w1, w2)
-    Q = orthogonalize(AOmega, orthogonalize_method)
-
-    AQ = weighted_eigen_plus_rank1_mm(L1, Q1, v2, Q, w1, w2)
-    QTAQ = Q.T @ AQ
-
-    W = (QTAQ + QTAQ.T) / 2.0
-
-    # compute new L and Q
-    try:
-        L_prime, S = torch_linalg.eigh(W, retry_float64=True)
-    except torch.linalg.LinAlgError:
-        return L1, Q1
-
-    L_prime, S = regularize_eigh(L=L_prime, Q=S, truncate=rank, tol=eig_tol, damping=damping, rdamping=rdamping)
-
-    if L_prime is None or S is None:
-        return L1, Q1
-
-    return L_prime, Q @ S
-
-
-# def adanystrom_update2(
-#     L1: torch.Tensor,
-#     Q1: torch.Tensor,
-#     v2: torch.Tensor,
-#     w1: float,
-#     w2: float,
-#     rank: int,
-# ):
-#     def A_mm(X):
-#         return weighted_eigen_plus_rank1_mm(L1=L1, Q1=Q1, v2=v2, B=X, w1=w1, w2=w2)
-
-#     return nystrom_approximation(A_mm, A_mm=A_mm, ndim=v2.numel(), rank=rank, device=L1.device, dtype=L1.dtype)
-
-class AdaNystrom(TensorTransform):
-    """Adagrad/RMSprop/Adam with Nyström-approximated covariance matrix.
-
-    Args:
-        rank (_type_): rank of Nyström approximation.
-        w1 (float, optional): weight of current covariance matrix. Defaults to 0.95.
-        w2 (float, optional): weight of new gradient in covariance matrix. Defaults to 0.05.
-        oversampling (int, optional): number of extra random vectors (top rank eigenvalues are kept). Defaults to 10.
-        eig_tol (float, optional):
-            removes eigenvalues this much smaller than largest eigenvalue when updating the preconditioner. Defaults to 1e-7.
-        damping (float, optional):
-            added to eigenvalues when updating the preconditioner. Defaults to 1e-8.
-        rdamping (float, optional):
-            added to eigenvalues when updating the preconditioner, relative to largest eigenvalue. Defaults to 0.
-        mm_tol (float, optional):
-            removes eigenvalues this much smaller than largest eigenvalue when computing the update. Defaults to 1e-7.
-        mm_truncate (int | None, optional):
-            uses top k eigenvalues to compute the update. Defaults to None.
-        mm_damping (float, optional):
-            added to eigenvalues when computing the update. Defaults to 1e-4.
-        mm_rdamping (float, optional):
-            added to eigenvalues when computing the update, relative to largest eigenvalue. Defaults to 0.
-        id_reg (float, optional):
-            multiplier to identity matrix added to preconditioner before computing update
-            If this value is given, solution from Nyström sketch-and-solve will be used to compute the update.
-            This value can't be too small (i.e. less than 1e-5) or the solver will be very unstable. Defaults to None.
-        concat_params (bool, optional):
-            whether to precondition all parameters at once if True, or each separately if False. Defaults to True.
-        update_freq (int, optional): update frequency. Defaults to 1.
-        inner (Chainable | None, optional): inner modules. Defaults to None.
-    """
-    def __init__(
-        self,
-        rank:int = 100,
-        beta=0.95,
-        oversampling: int = 10,
-        eig_tol: float | None = 1e-32,
-        damping: float = 0,
-        rdamping: float = 0,
-        mm_tol: float = 0,
-        mm_truncate: int | None = None,
-        mm_damping: float = 0,
-        mm_rdamping: float = 0,
-        id_reg: float | None = None,
-        orthogonalize_method: OrthogonalizeMethod = 'qr',
-        eigenbasis_optimizer: LREOptimizerBase | None = None,
-        orthogonalize_interval: int | None = 100,
-
-        concat_params: bool = True,
-        update_freq: int = 1,
-        inner: Chainable | None = None,
-    ):
-        defaults = locals().copy()
-        for k in ["self", "concat_params", "inner", "update_freq"]:
-            del defaults[k]
-
-        super().__init__(defaults, concat_params=concat_params, inner=inner, update_freq=update_freq)
-
-    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
-        state["step"] = state.get("step", 0) + 1
-        rank = setting["rank"]
-        device = tensor.device
-        dtype = tensor.dtype
-        beta = setting["beta"]
-
-        try:
-            if "L" not in state:
-                # use just tensor and zero L and Q with zero weight
-
-                L, Q = adanystrom_update(
-                    L1=torch.zeros(rank, device=device, dtype=dtype),
-                    Q1=torch.zeros((tensor.numel(), rank), device=device, dtype=dtype),
-                    v2=tensor.ravel(),
-                    w1=0,
-                    w2=1-beta,
-                    rank=rank,
-                    oversampling_p=setting["oversampling"],
-                    eig_tol=setting["eig_tol"],
-                    damping=setting["damping"],
-                    rdamping=setting["rdamping"],
-                    orthogonalize_method=setting["orthogonalize_method"],
-                )
-
-                state["L"] = state["L_reg"] = L
-                state["Q"] = state["Q_reg"] = Q
-
-            else:
-                L = state["L"]
-                Q = state["Q"]
-
-                w1 = beta
-                w2 = 1 - w1
-
-                # compute new factors (this function truncates them)
-                L_new, Q_new = adanystrom_update(
-                    L1=L,
-                    Q1=Q,
-                    v2=tensor.ravel(),
-                    w1=w1,
-                    w2=w2,
-                    rank=rank,
-                    oversampling_p=setting["oversampling"],
-                    eig_tol=setting["eig_tol"],
-                    damping=setting["damping"],
-                    rdamping=setting["rdamping"],
-                    orthogonalize_method=setting["orthogonalize_method"],
-                )
-
-                _eigengrad_update_state_(state=state, setting=setting, L_new=L_new, Q_new=Q_new)
-
-        except torch.linalg.LinAlgError:
-            pass
-
-    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
-        if "L_reg" not in state:
-            return tensor.clip(-0.1, 0.1)
-
-        if "eigenbasis_state" not in state:
-            state["eigenbasis_state"] = {}
-
-        return eigengrad_apply(
-            tensor=tensor,
-            L_reg = state["L_reg"],
-            Q_reg = state["Q_reg"],
-            beta = setting["beta"],
-            step = state["step"],
-            debias = True,
-            id_reg = setting["id_reg"],
-            eigenbasis_optimizer = setting["eigenbasis_optimizer"],
-            eigenbasis_state = state["eigenbasis_state"]
-        )

torchzero/modules/experimental/common_directions_whiten.py

@@ -1,142 +0,0 @@
-from collections import deque
-from typing import Literal
-
-import torch
-
-from torchzero.core import Chainable, TensorTransform
-from torchzero.linalg import matrix_power_eigh, torch_linalg, orthogonalize, OrthogonalizeMethod, regularize_eigh
-from torchzero.utils import TensorList, vec_to_tensors_
-
-
-def update_subspace_preconditioner_(
-    grad: torch.Tensor, # store grads and basis as vectors for matmul
-    basis: torch.Tensor, # ndim, k
-    accumulator_: torch.Tensor, # k, k
-    beta: float | None,
-):
-    projected = basis.T @ grad # k
-    outer = torch.outer(projected, projected)
-
-    if beta is None: accumulator_.add_(outer)
-    else: accumulator_.lerp_(outer, 1-beta)
-
-# yeah so I can also run subspace opts in this basis
-def apply_subspace_preconditioner(
-    tensor: torch.Tensor,
-    basis: torch.Tensor, # ndim, k
-    accumulator: torch.Tensor,
-    tol: float,
-    truncate: int | None,
-    damping: float,
-    rdamping: float,
-):
-    L, Q = torch_linalg.eigh(accumulator, retry_float64=True)
-    L, Q = regularize_eigh(L=L, Q=Q, truncate=truncate, tol=tol, damping=damping, rdamping=rdamping)
-
-    if L is None or Q is None:
-        return tensor.clip(-0.1, 0.1)
-
-    preconditioner = (Q * L.rsqrt().unsqueeze(-2)) @ Q.mH
-
-    tensor_projected = basis.T @ tensor # k
-    update_projected = preconditioner @ tensor_projected # k
-    return basis @ update_projected # d
-
-
-class CommonDirectionsWhiten(TensorTransform):
-    """Whitens in subspace spanned by history of gradient differences.
-
-    Args:
-        beta - for preconditioner itself in the basis.
-        basis_beta - how much basis is allowed to change.
-    """
-
-    def __init__(
-        self,
-        k: int = 100,
-        beta: float | None = 0.95,
-        basis_beta=0.95,
-        tol: float = 1e-7,
-        truncate: int | None = None,
-        damping: float = 1e-4,
-        rdamping: float = 0,
-        basis_type: Literal["gradients", "differences"] = "differences",
-        orthogonalize_method: OrthogonalizeMethod | None = 'newtonschulz',
-
-        concat_params: bool = True,
-        inner: Chainable | None = None,
-    ):
-        defaults = locals().copy()
-        for key in ["self", "inner", "concat_params"]:
-            del defaults[key]
-
-        super().__init__(defaults, concat_params=concat_params, inner=inner)
-
-    @torch.no_grad
-    def single_tensor_update(self, tensor, param, grad, loss, state, setting):
-        g = tensor.ravel()
-        k = setting['k']
-        beta = setting['beta']
-        basis_beta = setting['basis_beta']
-        step = state.get("step", 0)
-        state["step"] = step + 1
-
-        # initialize history
-        if 'history' not in state:
-            state['history'] = deque(maxlen=k)
-            state['accumulator'] = torch.eye(k, device=g.device, dtype=g.dtype)
-            state['basis'] = torch.zeros(g.numel(), k, device=g.device, dtype=g.dtype)
-
-        history: deque = state['history']
-        accumulator = state['accumulator']
-        basis = state['basis']
-        history.append(g)
-
-        # stack history to new basis term, if history isn't full, fill with random vecs
-        if len(history) < k:
-            basis_t = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
-            history_basis = torch.stack(tuple(history), -1)
-            basis_t[:, -len(history):] = history_basis
-
-        else:
-            basis_t = torch.stack(tuple(history), -1)
-
-        # in this case basis uses differences in gradients except last entry is the gradient
-        if setting["basis_type"] == "differences":
-            basis_t[:,:-1] = basis_t[:, :-1] - basis_t[:, 1:]
-
-        # normalize or orthonormalize new basis term
-        if setting["orthogonalize_method"] is not None:
-            basis_t = orthogonalize(basis_t, method = setting["orthogonalize_method"])
-        else:
-            basis_t = (basis_t - basis_t.mean()) / basis_t.std().clip(min=torch.finfo(g.dtype).tiny * 2)
-
-        # lerp basis
-        basis.lerp_(basis_t, 1-basis_beta)
-        basis = basis /  (1 - basis_beta ** (step+1)) # correct bias on basis EMA
-        update_subspace_preconditioner_(g, basis, accumulator, beta)
-
-    @torch.no_grad
-    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
-        g = tensor.ravel()
-
-        basis = state['basis']
-        accumulator = state['accumulator']
-        step = state["step"]
-        accumulator = accumulator / (1 - setting["beta"] ** (step+1)) # correct bias on accumulator EMA
-
-        try:
-            preconditioned = apply_subspace_preconditioner(
-                g,
-                basis,
-                accumulator,
-                tol=setting["tol"],
-                truncate=setting["truncate"],
-                damping=setting["damping"],
-                rdamping=setting["rdamping"],
-            )
-        except torch.linalg.LinAlgError:
-            preconditioned = g.clip(-0.1, 0.1)
-
-        return preconditioned.view_as(tensor)
-