torchax 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of torchax might be problematic. Click here for more details.

Files changed (32) hide show
  1. torchax/CONTRIBUTING.md +2 -2
  2. torchax/__init__.py +26 -24
  3. torchax/amp.py +332 -0
  4. torchax/config.py +25 -14
  5. torchax/configuration.py +30 -0
  6. torchax/decompositions.py +663 -195
  7. torchax/device_module.py +14 -1
  8. torchax/environment.py +0 -1
  9. torchax/export.py +26 -17
  10. torchax/flax.py +39 -0
  11. torchax/interop.py +288 -141
  12. torchax/mesh_util.py +220 -0
  13. torchax/ops/jaten.py +1723 -1297
  14. torchax/ops/jax_reimplement.py +23 -21
  15. torchax/ops/jc10d.py +5 -4
  16. torchax/ops/jimage.py +113 -0
  17. torchax/ops/jlibrary.py +9 -2
  18. torchax/ops/jtorch.py +237 -88
  19. torchax/ops/jtorchvision_nms.py +32 -43
  20. torchax/ops/mappings.py +77 -35
  21. torchax/ops/op_base.py +59 -32
  22. torchax/ops/ops_registry.py +40 -35
  23. torchax/tensor.py +442 -288
  24. torchax/train.py +38 -41
  25. torchax/util.py +88 -0
  26. torchax/view.py +377 -0
  27. {torchax-0.0.4.dist-info → torchax-0.0.6.dist-info}/METADATA +111 -145
  28. torchax-0.0.6.dist-info/RECORD +33 -0
  29. torchax/distributed.py +0 -246
  30. torchax-0.0.4.dist-info/RECORD +0 -27
  31. {torchax-0.0.4.dist-info → torchax-0.0.6.dist-info}/WHEEL +0 -0
  32. {torchax-0.0.4.dist-info → torchax-0.0.6.dist-info}/licenses/LICENSE +0 -0
torchax-0.0.4.dist-info/RECORD DELETED
@@ -1,27 +0,0 @@
1
- torchax/CONTRIBUTING.md,sha256=T6Uy7vZvzXArbPqqUswEtQBgTU_eRwtJnHOiZNth2Pw,1371
2
- torchax/__init__.py,sha256=0NoxpzaPIMYMJXZz_JiRCTYZCuICnTFrLDm5cZKGIm0,3242
3
- torchax/config.py,sha256=Zpfcn7Q3QsKnn4cCx1-bYxA8q_lqrWJcUz30B3QcIc4,535
4
- torchax/decompositions.py,sha256=C067xSd9MpGs1SDqlC5T9RnqNj-i-m31mysfFGDxLEQ,11221
5
- torchax/device_module.py,sha256=NNEGSfk9ApVVSy5dRjHkvGuNGamopCEjSlItcNlAIbE,253
6
- torchax/distributed.py,sha256=CXkpV0K2Oas7fcRA-i3s2VXTz8l4CEI7_NnbYOsHtRw,7519
7
- torchax/environment.py,sha256=daEdpEyAJIa8b2VkCqSKcw8PaExcB6Qro80XNes_sHA,2
8
- torchax/export.py,sha256=MG19Y0QcRcX3S3gruy9K-wFGdemD4I0ii03HWQM31xk,8853
9
- torchax/interop.py,sha256=fncREqC1Nib-ud7wyekmtKaif7GSDql7ITmzTt5A9e4,6774
10
- torchax/tensor.py,sha256=89yXmyO1uT6F1KI-Kn6NCnIieCiiMsDutgS071_9p64,16644
11
- torchax/tf_integration.py,sha256=d_h4vSJm7N9rJXpUPNCDOiUz3J1-UPo3KU8D9Wi4nnc,4074
12
- torchax/train.py,sha256=Ym_SC0WOmhQMhSCH0aZ-zgP7Q19WdCDTeoI-JBFrqx8,3930
13
- torchax/types.py,sha256=j4ERjkgDgwhgi9zrwwbbiv4HMDlrJ1IEMUCmP_BIJ9M,388
14
- torchax/ops/__init__.py,sha256=Vr1p8zDHwfXZBUbw70iNiCJLZLNdI6gR_vUlaiA7Usg,270
15
- torchax/ops/jaten.py,sha256=ycTSyCzgJpz5zu936BTZXmRJWuzhxeGOOjoffU00c3g,154999
16
- torchax/ops/jax_reimplement.py,sha256=29SGPF0bszgCuyVxiZ4A2xTPtAuPXBDvDG79M3C0Vxo,7468
17
- torchax/ops/jc10d.py,sha256=DiJjiUFHyYmFUHEgqEDbGsLaCsJs94xJlBatRwXiPfg,1317
18
- torchax/ops/jlibrary.py,sha256=R22GUlK6lUivuDSmvtqwbV9ku_uT0ddLcQcwHI96vqM,2927
19
- torchax/ops/jtorch.py,sha256=fWU2z5uD3FpVE5ui-O5nfHEule2hod4BqCnsbV5jVfQ,14463
20
- torchax/ops/jtorchvision_nms.py,sha256=SiMN9bUV7GADD2S2PkH9spn5oQGgr2jhFQ1nyd3eHgI,8873
21
- torchax/ops/mappings.py,sha256=a8C5CYIEucGIxUPQHCsDZypBC3b_BYtOKQTWf9bdiRc,3060
22
- torchax/ops/op_base.py,sha256=EdlAx3oqU2brMngjZvTxECin0eNiPdAUfOLRtFL3tgg,2900
23
- torchax/ops/ops_registry.py,sha256=nXCAXiQNpC0uBUl8xEim7TOUTyROZm_4H1LaOvYwJVQ,1241
24
- torchax-0.0.4.dist-info/METADATA,sha256=4VQb39quiN9U9j955fs9zY00x4t4Orubo92ONCCLyug,11119
25
- torchax-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
26
- torchax-0.0.4.dist-info/licenses/LICENSE,sha256=ZHyir3-ltOerFLt9JH1bjf7lIxIWipFmqeMnB_8z_aU,1498
27
- torchax-0.0.4.dist-info/RECORD,,