haliax 1.4.dev366__tar.gz → 1.4.dev367__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 (112) hide show
  1. {haliax-1.4.dev366 → haliax-1.4.dev367}/AGENTS.md +3 -9
  2. {haliax-1.4.dev366 → haliax-1.4.dev367}/PKG-INFO +1 -15
  3. {haliax-1.4.dev366 → haliax-1.4.dev367}/pyproject.toml +11 -4
  4. haliax-1.4.dev367/src/haliax/__about__.py +1 -0
  5. haliax-1.4.dev366/src/haliax/__about__.py +0 -1
  6. {haliax-1.4.dev366 → haliax-1.4.dev367}/.coveragerc +0 -0
  7. {haliax-1.4.dev366 → haliax-1.4.dev367}/.flake8 +0 -0
  8. {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/publish_dev.yaml +0 -0
  9. {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/run_pre_commit.yaml +0 -0
  10. {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
  11. {haliax-1.4.dev366 → haliax-1.4.dev367}/.github/workflows/run_tests.yaml +0 -0
  12. {haliax-1.4.dev366 → haliax-1.4.dev367}/.gitignore +0 -0
  13. {haliax-1.4.dev366 → haliax-1.4.dev367}/.pre-commit-config.yaml +0 -0
  14. {haliax-1.4.dev366 → haliax-1.4.dev367}/.readthedocs.yaml +0 -0
  15. {haliax-1.4.dev366 → haliax-1.4.dev367}/CONTRIBUTING.md +0 -0
  16. {haliax-1.4.dev366 → haliax-1.4.dev367}/LICENSE +0 -0
  17. {haliax-1.4.dev366 → haliax-1.4.dev367}/README.md +0 -0
  18. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/api.md +0 -0
  19. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/broadcasting.md +0 -0
  20. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/cheatsheet.md +0 -0
  21. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/css/material.css +0 -0
  22. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/css/mkdocstrings.css +0 -0
  23. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/faq.md +0 -0
  24. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/data_parallel_mesh.png +0 -0
  25. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/data_parallel_mesh_replicated.png +0 -0
  26. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_1d.png +0 -0
  27. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_1d_zero.png +0 -0
  28. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d.png +0 -0
  29. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
  30. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
  31. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
  32. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
  33. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/figures/device_mesh_2d_zero.png +0 -0
  34. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/fp8.md +0 -0
  35. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/index.md +0 -0
  36. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/indexing.md +0 -0
  37. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/matmul.md +0 -0
  38. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/nn.md +0 -0
  39. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/partitioning.md +0 -0
  40. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/rearrange.ipynb +0 -0
  41. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/rearrange.md +0 -0
  42. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/requirements.txt +0 -0
  43. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/scan.md +0 -0
  44. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/state-dict.md +0 -0
  45. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/tutorial.md +0 -0
  46. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/typing.md +0 -0
  47. {haliax-1.4.dev366 → haliax-1.4.dev367}/docs/vmap.md +0 -0
  48. {haliax-1.4.dev366 → haliax-1.4.dev367}/mkdocs.yml +0 -0
  49. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/__init__.py +0 -0
  50. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/__init__.py +0 -0
  51. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/compile_utils.py +0 -0
  52. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/dot.py +0 -0
  53. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/einsum.py +0 -0
  54. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/fp8.py +0 -0
  55. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/parsing.py +0 -0
  56. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/rearrange.py +0 -0
  57. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/scan.py +0 -0
  58. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/state_dict.py +0 -0
  59. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/_src/util.py +0 -0
  60. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/axis.py +0 -0
  61. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/core.py +0 -0
  62. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/debug.py +0 -0
  63. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/haxtyping.py +0 -0
  64. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/hof.py +0 -0
  65. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/jax_utils.py +0 -0
  66. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/__init__.py +0 -0
  67. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/activations.py +0 -0
  68. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/attention.py +0 -0
  69. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/conv.py +0 -0
  70. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/dropout.py +0 -0
  71. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/embedding.py +0 -0
  72. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/linear.py +0 -0
  73. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/loss.py +0 -0
  74. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/mlp.py +0 -0
  75. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/normalization.py +0 -0
  76. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/pool.py +0 -0
  77. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/nn/scan.py +0 -0
  78. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/ops.py +0 -0
  79. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/partitioning.py +0 -0
  80. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/quantization.py +0 -0
  81. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/random.py +0 -0
  82. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/specialized_fns.py +0 -0
  83. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/state_dict.py +0 -0
  84. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/tree_util.py +0 -0
  85. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/types.py +0 -0
  86. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/util.py +0 -0
  87. {haliax-1.4.dev366 → haliax-1.4.dev367}/src/haliax/wrap.py +0 -0
  88. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/core_test.py +0 -0
  89. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_attention.py +0 -0
  90. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_axis.py +0 -0
  91. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_conv.py +0 -0
  92. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_debug.py +0 -0
  93. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_dot.py +0 -0
  94. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_dtype_typing.py +0 -0
  95. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_einsum.py +0 -0
  96. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_fp8.py +0 -0
  97. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_hof.py +0 -0
  98. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_int8.py +0 -0
  99. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_namedarray_typing.py +0 -0
  100. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_nn.py +0 -0
  101. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_ops.py +0 -0
  102. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_parsing.py +0 -0
  103. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_partitioning.py +0 -0
  104. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_pool.py +0 -0
  105. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_random.py +0 -0
  106. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_rearrange.py +0 -0
  107. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_scan.py +0 -0
  108. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_scatter_gather.py +0 -0
  109. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_specialized_fns.py +0 -0
  110. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_state_dict.py +0 -0
  111. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_tree_util.py +0 -0
  112. {haliax-1.4.dev366 → haliax-1.4.dev367}/tests/test_utils.py +0 -0
@@ -25,7 +25,8 @@ repository. Follow these notes when implementing new features or fixing bugs.
25
25
  * **Formatting and Linting**: We use `ruff` via `pre-commit`.
26
26
  * **Typing**: the code base uses `mypy` for static type checking. `mypy` is run by pre‑commit and the
27
27
  configuration is found in `pyproject.toml`.
28
- * **Run `pre-commit run --all-files`** before committing. The CI workflows run the same checks.
28
+ * **Run `uv run pre-commit run --all-files`** before committing. The CI workflows run the same checks.
29
+ * **Use `uv run` for commands.** When running tools like `pytest` or other scripts, invoke them via `uv run` so the development dependencies are active.
29
30
  * **Doc Strings**: All public functions, classes, and modules should have docstrings, unless
30
31
  their purpose is painfully obvious. Use
31
32
  [Google style](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) for
@@ -44,12 +45,10 @@ repository. Follow these notes when implementing new features or fixing bugs.
44
45
 
45
46
  ## Testing
46
47
 
47
- * Tests are executed with `pytest`. The default workflow runs
48
- `pytest tests -m "not entry and not slow and not ray"`.
48
+ * Tests are executed with `pytest`. The default workflow runs `uv run pytest tests`.
49
49
  * In general, never relax tolerances in floating point tests unless specifically discussed with the
50
50
  team. Use `assert_allclose` with appropriate tolerances for numerical comparisons. We typically use
51
51
  1e-4 for more complex modules, and 1e-5 for simpler ones.
52
- * Tests should be reasonably fast. Mark long-running tests with @pytest.mark.slow so they are excluded from the default suite.
53
52
  * Always mark tests that depend on pytorch with `@skip_if_no_torch` to ensure they are skipped
54
53
  when PyTorch is not available. This is particularly important for tests that require PyTorch-specific
55
54
  functionality.
@@ -59,11 +58,6 @@ repository. Follow these notes when implementing new features or fixing bugs.
59
58
 
60
59
  * **Generic code**: many utilities are written with Python generics and dataclasses. Where possible,
61
60
  write reusable functions or classes that operate over TypeVars instead of hard coding concrete types.
62
- * **Configurations**: configuration files are dataclasses loaded via `draccus`. Keep configs
63
- declarative and typed.
64
- * **Reproducibility**: Levanter aims for deterministic training where possible. Avoid sources of
65
- nondeterminism unless explicitly required.
66
- * Prefer Stacked with fold or scan over writing custom loops, for better compile times and gradient checkpointing support
67
61
 
68
62
  ## Library conventions
69
63
  - Haliax revolves around `NamedArray` and explicit `Axis` objects. Prefer APIs that accept
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: haliax
3
- Version: 1.4.dev366
3
+ Version: 1.4.dev367
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/
@@ -19,20 +19,6 @@ Requires-Dist: equinox>=0.10.6
19
19
  Requires-Dist: jaxtyping>=0.2.20
20
20
  Requires-Dist: jmp>=0.0.4
21
21
  Requires-Dist: safetensors>=0.4.3
22
- Provides-Extra: dev
23
- Requires-Dist: chex>=0.1.86; extra == 'dev'
24
- Requires-Dist: mkdocs-include-markdown-plugin; extra == 'dev'
25
- Requires-Dist: mkdocs-literate-nav>=0.6.0; extra == 'dev'
26
- Requires-Dist: mkdocs-macros-plugin>=0.7.0; extra == 'dev'
27
- Requires-Dist: mkdocs-material>=7.3.3; extra == 'dev'
28
- Requires-Dist: mkdocs>=1.4.3; extra == 'dev'
29
- Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'dev'
30
- Requires-Dist: mkdocstrings>=0.22.0; extra == 'dev'
31
- Requires-Dist: mypy>=0.910; extra == 'dev'
32
- Requires-Dist: pre-commit; extra == 'dev'
33
- Requires-Dist: pygments; extra == 'dev'
34
- Requires-Dist: pymdown-extensions; extra == 'dev'
35
- Requires-Dist: pytest>=7.4.0; extra == 'dev'
36
22
  Description-Content-Type: text/markdown
37
23
 
38
24
  <!--haliax-intro-start-->
@@ -31,13 +31,20 @@ dependencies = [
31
31
  ]
32
32
  dynamic =[ "version" ]
33
33
 
34
- [project.optional-dependencies]
35
- dev=["pytest >= 7.4.0", "mypy >= 0.910", "mkdocs >= 1.4.3", "mkdocs-material >= 7.3.3", "mkdocstrings >= 0.22.0",
36
- "mkdocs-literate-nav >= 0.6.0", "mkdocs-macros-plugin >= 0.7.0", "mkdocstrings-python >= 1.1.2",
34
+
35
+ [dependency-groups]
36
+ dev = [
37
+ "pytest >= 7.4.0",
38
+ "mypy >= 0.910",
39
+ "mkdocs >= 1.4.3",
40
+ "mkdocs-material >= 7.3.3",
41
+ "mkdocstrings >= 0.22.0",
42
+ "mkdocs-literate-nav >= 0.6.0",
43
+ "mkdocs-macros-plugin >= 0.7.0",
44
+ "mkdocstrings-python >= 1.1.2",
37
45
  "mkdocs-include-markdown-plugin",
38
46
  "pymdown-extensions",
39
47
  "pygments",
40
- "pymdown-extensions",
41
48
  "chex>=0.1.86",
42
49
  "pre-commit",
43
50
  ]
@@ -0,0 +1 @@
1
+ __version__ = "1.4.dev367"
@@ -1 +0,0 @@
1
- __version__ = "1.4.dev366"
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes