haliax 1.4.dev439__tar.gz → 1.4.dev441__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 (134) hide show
  1. {haliax-1.4.dev439 → haliax-1.4.dev441}/PKG-INFO +1 -1
  2. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/__about__.py +1 -1
  3. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/mup.py +1 -1
  4. {haliax-1.4.dev439 → haliax-1.4.dev441}/.agents/projects/api_parity.md +0 -0
  5. {haliax-1.4.dev439 → haliax-1.4.dev441}/.coveragerc +0 -0
  6. {haliax-1.4.dev439 → haliax-1.4.dev441}/.flake8 +0 -0
  7. {haliax-1.4.dev439 → haliax-1.4.dev441}/.github/workflows/publish_dev.yaml +0 -0
  8. {haliax-1.4.dev439 → haliax-1.4.dev441}/.github/workflows/run_pre_commit.yaml +0 -0
  9. {haliax-1.4.dev439 → haliax-1.4.dev441}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
  10. {haliax-1.4.dev439 → haliax-1.4.dev441}/.github/workflows/run_tests.yaml +0 -0
  11. {haliax-1.4.dev439 → haliax-1.4.dev441}/.gitignore +0 -0
  12. {haliax-1.4.dev439 → haliax-1.4.dev441}/.playbooks/add-types.md +0 -0
  13. {haliax-1.4.dev439 → haliax-1.4.dev441}/.playbooks/wrap-non-named.md +0 -0
  14. {haliax-1.4.dev439 → haliax-1.4.dev441}/.pre-commit-config.yaml +0 -0
  15. {haliax-1.4.dev439 → haliax-1.4.dev441}/.readthedocs.yaml +0 -0
  16. {haliax-1.4.dev439 → haliax-1.4.dev441}/AGENTS.md +0 -0
  17. {haliax-1.4.dev439 → haliax-1.4.dev441}/AUTHORS.md +0 -0
  18. {haliax-1.4.dev439 → haliax-1.4.dev441}/CONTRIBUTING.md +0 -0
  19. {haliax-1.4.dev439 → haliax-1.4.dev441}/CONTRIBUTORS.md +0 -0
  20. {haliax-1.4.dev439 → haliax-1.4.dev441}/LICENSE +0 -0
  21. {haliax-1.4.dev439 → haliax-1.4.dev441}/README.md +0 -0
  22. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/api.md +0 -0
  23. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/broadcasting.md +0 -0
  24. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/cheatsheet.md +0 -0
  25. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/css/material.css +0 -0
  26. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/css/mkdocstrings.css +0 -0
  27. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/faq.md +0 -0
  28. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/data_parallel_mesh.png +0 -0
  29. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/data_parallel_mesh_replicated.png +0 -0
  30. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_1d.png +0 -0
  31. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_1d_zero.png +0 -0
  32. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_2d.png +0 -0
  33. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
  34. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
  35. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
  36. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
  37. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/figures/device_mesh_2d_zero.png +0 -0
  38. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/fp8.md +0 -0
  39. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/index.md +0 -0
  40. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/indexing.md +0 -0
  41. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/matmul.md +0 -0
  42. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/nn.md +0 -0
  43. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/partitioning.md +0 -0
  44. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/primer.md +0 -0
  45. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/rearrange.ipynb +0 -0
  46. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/rearrange.md +0 -0
  47. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/requirements.txt +0 -0
  48. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/scan.md +0 -0
  49. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/state-dict.md +0 -0
  50. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/tutorial.md +0 -0
  51. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/typing.md +0 -0
  52. {haliax-1.4.dev439 → haliax-1.4.dev441}/docs/vmap.md +0 -0
  53. {haliax-1.4.dev439 → haliax-1.4.dev441}/etc/license_header.txt +0 -0
  54. {haliax-1.4.dev439 → haliax-1.4.dev441}/mkdocs.yml +0 -0
  55. {haliax-1.4.dev439 → haliax-1.4.dev441}/pyproject.toml +0 -0
  56. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/__init__.py +0 -0
  57. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/__init__.py +0 -0
  58. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/compile_utils.py +0 -0
  59. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/dot.py +0 -0
  60. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/einsum.py +0 -0
  61. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/fp8.py +0 -0
  62. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/parsing.py +0 -0
  63. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/rearrange.py +0 -0
  64. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/scan.py +0 -0
  65. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/state_dict.py +0 -0
  66. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/_src/util.py +0 -0
  67. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/axis.py +0 -0
  68. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/core.py +0 -0
  69. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/debug.py +0 -0
  70. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/fft.py +0 -0
  71. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/field.py +0 -0
  72. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/haxtyping.py +0 -0
  73. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/hof.py +0 -0
  74. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/jax_utils.py +0 -0
  75. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/__init__.py +0 -0
  76. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/activations.py +0 -0
  77. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/attention.py +0 -0
  78. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/conv.py +0 -0
  79. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/dropout.py +0 -0
  80. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/embedding.py +0 -0
  81. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/linear.py +0 -0
  82. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/loss.py +0 -0
  83. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/mlp.py +0 -0
  84. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/normalization.py +0 -0
  85. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/pool.py +0 -0
  86. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/scan.py +0 -0
  87. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/ops.py +0 -0
  88. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/partitioning.py +0 -0
  89. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/poly.py +0 -0
  90. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/quantization.py +0 -0
  91. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/random.py +0 -0
  92. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/specialized_fns.py +0 -0
  93. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/state_dict.py +0 -0
  94. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/tree.py +0 -0
  95. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/tree_util.py +0 -0
  96. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/types.py +0 -0
  97. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/util.py +0 -0
  98. {haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/wrap.py +0 -0
  99. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/core_test.py +0 -0
  100. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_attention.py +0 -0
  101. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_axis.py +0 -0
  102. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_bitwise_ops.py +0 -0
  103. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_conv.py +0 -0
  104. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_debug.py +0 -0
  105. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_dot.py +0 -0
  106. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_dtype_typing.py +0 -0
  107. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_einsum.py +0 -0
  108. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_fft.py +0 -0
  109. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_field.py +0 -0
  110. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_fp8.py +0 -0
  111. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_hof.py +0 -0
  112. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_int8.py +0 -0
  113. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_moe_linear.py +0 -0
  114. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_mup_coordinate_check.py +0 -0
  115. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_mup_embedding.py +0 -0
  116. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_mup_linear.py +0 -0
  117. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_namedarray_typing.py +0 -0
  118. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_nan_reductions.py +0 -0
  119. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_nn.py +0 -0
  120. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_ops.py +0 -0
  121. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_parsing.py +0 -0
  122. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_partitioning.py +0 -0
  123. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_poly_ops.py +0 -0
  124. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_pool.py +0 -0
  125. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_random.py +0 -0
  126. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_rearrange.py +0 -0
  127. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_scan.py +0 -0
  128. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_scatter_gather.py +0 -0
  129. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_specialized_fns.py +0 -0
  130. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_state_dict.py +0 -0
  131. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_tree_util.py +0 -0
  132. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_utils.py +0 -0
  133. {haliax-1.4.dev439 → haliax-1.4.dev441}/tests/test_visualize_sharding.py +0 -0
  134. {haliax-1.4.dev439 → haliax-1.4.dev441}/uv.lock +0 -0
{haliax-1.4.dev439 → haliax-1.4.dev441}/PKG-INFO RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: haliax
3
- Version: 1.4.dev439
3
+ Version: 1.4.dev441
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/
{haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/__about__.py RENAMED
@@ -3,4 +3,4 @@
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
 
6
- __version__ = "1.4.dev439"
6
+ __version__ = "1.4.dev441"
{haliax-1.4.dev439 → haliax-1.4.dev441}/src/haliax/nn/mup.py RENAMED
@@ -168,7 +168,7 @@ class EmbeddingStandardParam(AbstractEmbeddingReparam):
168
168
 
169
169
  @property
170
170
  def unembed_active_scale(self):
171
- return 1 / hax.axis_size(self.Embed)
171
+ return 1
172
172
 
173
173
 
174
174
  class EmbeddingMup(AbstractEmbeddingReparam):
{haliax-1.4.dev439 → haliax-1.4.dev441}/.coveragerc RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/.flake8 RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/.gitignore RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/AGENTS.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/AUTHORS.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/CONTRIBUTING.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/CONTRIBUTORS.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/LICENSE RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/README.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/api.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/faq.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/fp8.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/index.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/matmul.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/nn.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/primer.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/scan.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/typing.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/docs/vmap.md RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/mkdocs.yml RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/pyproject.toml RENAMED
File without changes
{haliax-1.4.dev439 → haliax-1.4.dev441}/uv.lock RENAMED
File without changes