haliax 1.4.dev330__tar.gz → 1.4.dev331__tar.gz

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.
Files changed (104) hide show
  1. {haliax-1.4.dev330 → haliax-1.4.dev331}/PKG-INFO +3 -3
  2. {haliax-1.4.dev330 → haliax-1.4.dev331}/README.md +2 -2
  3. haliax-1.4.dev331/src/haliax/__about__.py +1 -0
  4. haliax-1.4.dev330/src/haliax/__about__.py +0 -1
  5. {haliax-1.4.dev330 → haliax-1.4.dev331}/.coveragerc +0 -0
  6. {haliax-1.4.dev330 → haliax-1.4.dev331}/.flake8 +0 -0
  7. {haliax-1.4.dev330 → haliax-1.4.dev331}/.github/workflows/publish_dev.yaml +0 -0
  8. {haliax-1.4.dev330 → haliax-1.4.dev331}/.github/workflows/run_pre_commit.yaml +0 -0
  9. {haliax-1.4.dev330 → haliax-1.4.dev331}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
  10. {haliax-1.4.dev330 → haliax-1.4.dev331}/.github/workflows/run_tests.yaml +0 -0
  11. {haliax-1.4.dev330 → haliax-1.4.dev331}/.gitignore +0 -0
  12. {haliax-1.4.dev330 → haliax-1.4.dev331}/.pre-commit-config.yaml +0 -0
  13. {haliax-1.4.dev330 → haliax-1.4.dev331}/.readthedocs.yaml +0 -0
  14. {haliax-1.4.dev330 → haliax-1.4.dev331}/CONTRIBUTING.md +0 -0
  15. {haliax-1.4.dev330 → haliax-1.4.dev331}/LICENSE +0 -0
  16. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/api.md +0 -0
  17. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/broadcasting.md +0 -0
  18. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/cheatsheet.md +0 -0
  19. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/css/material.css +0 -0
  20. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/css/mkdocstrings.css +0 -0
  21. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/faq.md +0 -0
  22. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/data_parallel_mesh.png +0 -0
  23. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/data_parallel_mesh_replicated.png +0 -0
  24. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_1d.png +0 -0
  25. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_1d_zero.png +0 -0
  26. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_2d.png +0 -0
  27. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
  28. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
  29. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
  30. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
  31. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/figures/device_mesh_2d_zero.png +0 -0
  32. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/fp8.md +0 -0
  33. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/hof.md +0 -0
  34. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/index.md +0 -0
  35. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/indexing.md +0 -0
  36. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/matmul.md +0 -0
  37. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/nn.md +0 -0
  38. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/partitioning.md +0 -0
  39. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/rearrange.ipynb +0 -0
  40. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/rearrange.md +0 -0
  41. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/requirements.txt +0 -0
  42. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/state-dict.md +0 -0
  43. {haliax-1.4.dev330 → haliax-1.4.dev331}/docs/tutorial.md +0 -0
  44. {haliax-1.4.dev330 → haliax-1.4.dev331}/mkdocs.yml +0 -0
  45. {haliax-1.4.dev330 → haliax-1.4.dev331}/pyproject.toml +0 -0
  46. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/__init__.py +0 -0
  47. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/__init__.py +0 -0
  48. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/compile_utils.py +0 -0
  49. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/dot.py +0 -0
  50. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/einsum.py +0 -0
  51. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/fp8.py +0 -0
  52. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/parsing.py +0 -0
  53. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/rearrange.py +0 -0
  54. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/state_dict.py +0 -0
  55. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/_src/util.py +0 -0
  56. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/axis.py +0 -0
  57. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/core.py +0 -0
  58. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/debug.py +0 -0
  59. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/hof.py +0 -0
  60. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/jax_utils.py +0 -0
  61. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/__init__.py +0 -0
  62. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/activations.py +0 -0
  63. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/attention.py +0 -0
  64. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/conv.py +0 -0
  65. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/dropout.py +0 -0
  66. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/embedding.py +0 -0
  67. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/linear.py +0 -0
  68. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/loss.py +0 -0
  69. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/mlp.py +0 -0
  70. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/normalization.py +0 -0
  71. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/pool.py +0 -0
  72. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/nn/scan.py +0 -0
  73. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/ops.py +0 -0
  74. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/partitioning.py +0 -0
  75. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/quantization.py +0 -0
  76. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/random.py +0 -0
  77. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/specialized_fns.py +0 -0
  78. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/state_dict.py +0 -0
  79. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/tree_util.py +0 -0
  80. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/types.py +0 -0
  81. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/util.py +0 -0
  82. {haliax-1.4.dev330 → haliax-1.4.dev331}/src/haliax/wrap.py +0 -0
  83. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/core_test.py +0 -0
  84. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_attention.py +0 -0
  85. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_axis.py +0 -0
  86. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_conv.py +0 -0
  87. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_debug.py +0 -0
  88. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_dot.py +0 -0
  89. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_einsum.py +0 -0
  90. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_fp8.py +0 -0
  91. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_hof.py +0 -0
  92. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_int8.py +0 -0
  93. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_nn.py +0 -0
  94. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_ops.py +0 -0
  95. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_parsing.py +0 -0
  96. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_partitioning.py +0 -0
  97. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_pool.py +0 -0
  98. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_random.py +0 -0
  99. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_rearrange.py +0 -0
  100. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_scan.py +0 -0
  101. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_specialized_fns.py +0 -0
  102. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_state_dict.py +0 -0
  103. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_tree_util.py +0 -0
  104. {haliax-1.4.dev330 → haliax-1.4.dev331}/tests/test_utils.py +0 -0
{haliax-1.4.dev330 → haliax-1.4.dev331}/PKG-INFO RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: haliax
3
- Version: 1.4.dev330
3
+ Version: 1.4.dev331
4
4
  Summary: Named Tensors for Legible Deep Learning in JAX
5
5
  Project-URL: Homepage, https://github.com/stanford-crfm/haliax
6
6
  Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/
@@ -61,8 +61,8 @@ Despite the focus on legibility, Haliax
61
61
  is also **fast**, typically about as fast as "pure" JAX code.
62
62
  Haliax is also built to be **scalable**: it
63
63
  can support [Fully-Sharded Data Parallelism (FSDP)](https://engineering.fb.com/2021/07/15/open-source/fsdp/) and Tensor Parallelism with [just a few lines of code](https://colab.research.google.com/drive/1QX4yH3zRFF3Xiibf1aahETcSQ5nbcUMz). Haliax powers [Levanter](https://github.com/stanford-crfm/levanter),
64
- our companion library for training large language models and other foundation models, with scale proven up to 20B parameters
65
- and up to a TPU v3-256 pod slice.
64
+ our companion library for training large language models and other foundation models, with scale proven up to 70B parameters
65
+ and up to TPU v4-2048.
66
66
 
67
67
  ## Example: Attention
68
68
 
{haliax-1.4.dev330 → haliax-1.4.dev331}/README.md RENAMED
@@ -25,8 +25,8 @@ Despite the focus on legibility, Haliax
25
25
  is also **fast**, typically about as fast as "pure" JAX code.
26
26
  Haliax is also built to be **scalable**: it
27
27
  can support [Fully-Sharded Data Parallelism (FSDP)](https://engineering.fb.com/2021/07/15/open-source/fsdp/) and Tensor Parallelism with [just a few lines of code](https://colab.research.google.com/drive/1QX4yH3zRFF3Xiibf1aahETcSQ5nbcUMz). Haliax powers [Levanter](https://github.com/stanford-crfm/levanter),
28
- our companion library for training large language models and other foundation models, with scale proven up to 20B parameters
29
- and up to a TPU v3-256 pod slice.
28
+ our companion library for training large language models and other foundation models, with scale proven up to 70B parameters
29
+ and up to TPU v4-2048.
30
30
 
31
31
  ## Example: Attention
32
32
 
haliax-1.4.dev331/src/haliax/__about__.py ADDED
@@ -0,0 +1 @@
1
+ __version__ = "1.4.dev331"
haliax-1.4.dev330/src/haliax/__about__.py DELETED
@@ -1 +0,0 @@
1
- __version__ = "1.4.dev330"
{haliax-1.4.dev330 → haliax-1.4.dev331}/.coveragerc RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/.flake8 RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/.gitignore RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/CONTRIBUTING.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/LICENSE RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/docs/api.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/docs/faq.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/docs/fp8.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/docs/hof.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/docs/index.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/docs/matmul.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/docs/nn.md RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/mkdocs.yml RENAMED
File without changes
{haliax-1.4.dev330 → haliax-1.4.dev331}/pyproject.toml RENAMED
File without changes