haliax 1.4.dev299__tar.gz → 1.4.dev300__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 (99) hide show
  1. {haliax-1.4.dev299 → haliax-1.4.dev300}/.github/workflows/run_pre_commit.yaml +1 -1
  2. haliax-1.4.dev300/.github/workflows/run_quick_levanter_tests.yaml +36 -0
  3. {haliax-1.4.dev299 → haliax-1.4.dev300}/.github/workflows/run_tests.yaml +3 -6
  4. {haliax-1.4.dev299 → haliax-1.4.dev300}/PKG-INFO +1 -1
  5. haliax-1.4.dev300/src/haliax/__about__.py +1 -0
  6. haliax-1.4.dev299/src/haliax/__about__.py +0 -1
  7. {haliax-1.4.dev299 → haliax-1.4.dev300}/.coveragerc +0 -0
  8. {haliax-1.4.dev299 → haliax-1.4.dev300}/.flake8 +0 -0
  9. {haliax-1.4.dev299 → haliax-1.4.dev300}/.github/workflows/publish_dev.yaml +0 -0
  10. {haliax-1.4.dev299 → haliax-1.4.dev300}/.gitignore +0 -0
  11. {haliax-1.4.dev299 → haliax-1.4.dev300}/.pre-commit-config.yaml +0 -0
  12. {haliax-1.4.dev299 → haliax-1.4.dev300}/.readthedocs.yaml +0 -0
  13. {haliax-1.4.dev299 → haliax-1.4.dev300}/CONTRIBUTING.md +0 -0
  14. {haliax-1.4.dev299 → haliax-1.4.dev300}/LICENSE +0 -0
  15. {haliax-1.4.dev299 → haliax-1.4.dev300}/README.md +0 -0
  16. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/api.md +0 -0
  17. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/broadcasting.md +0 -0
  18. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/cheatsheet.md +0 -0
  19. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/css/material.css +0 -0
  20. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/css/mkdocstrings.css +0 -0
  21. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/faq.md +0 -0
  22. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/data_parallel_mesh.png +0 -0
  23. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/data_parallel_mesh_replicated.png +0 -0
  24. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_1d.png +0 -0
  25. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_1d_zero.png +0 -0
  26. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d.png +0 -0
  27. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
  28. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
  29. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
  30. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
  31. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/figures/device_mesh_2d_zero.png +0 -0
  32. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/fp8.md +0 -0
  33. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/hof.md +0 -0
  34. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/index.md +0 -0
  35. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/indexing.md +0 -0
  36. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/matmul.md +0 -0
  37. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/nn.md +0 -0
  38. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/partitioning.md +0 -0
  39. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/rearrange.ipynb +0 -0
  40. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/rearrange.md +0 -0
  41. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/requirements.txt +0 -0
  42. {haliax-1.4.dev299 → haliax-1.4.dev300}/docs/tutorial.md +0 -0
  43. {haliax-1.4.dev299 → haliax-1.4.dev300}/mkdocs.yml +0 -0
  44. {haliax-1.4.dev299 → haliax-1.4.dev300}/pyproject.toml +0 -0
  45. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/__init__.py +0 -0
  46. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/__init__.py +0 -0
  47. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/compile_utils.py +0 -0
  48. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/dot.py +0 -0
  49. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/einsum.py +0 -0
  50. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/fp8.py +0 -0
  51. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/parsing.py +0 -0
  52. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/rearrange.py +0 -0
  53. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/_src/util.py +0 -0
  54. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/axis.py +0 -0
  55. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/core.py +0 -0
  56. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/debug.py +0 -0
  57. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/hof.py +0 -0
  58. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/jax_utils.py +0 -0
  59. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/__init__.py +0 -0
  60. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/activations.py +0 -0
  61. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/attention.py +0 -0
  62. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/conv.py +0 -0
  63. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/dropout.py +0 -0
  64. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/embedding.py +0 -0
  65. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/linear.py +0 -0
  66. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/loss.py +0 -0
  67. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/mlp.py +0 -0
  68. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/normalization.py +0 -0
  69. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/pool.py +0 -0
  70. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/nn/scan.py +0 -0
  71. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/ops.py +0 -0
  72. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/partitioning.py +0 -0
  73. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/quantization.py +0 -0
  74. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/random.py +0 -0
  75. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/specialized_fns.py +0 -0
  76. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/tree_util.py +0 -0
  77. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/types.py +0 -0
  78. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/util.py +0 -0
  79. {haliax-1.4.dev299 → haliax-1.4.dev300}/src/haliax/wrap.py +0 -0
  80. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/core_test.py +0 -0
  81. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_attention.py +0 -0
  82. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_axis.py +0 -0
  83. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_conv.py +0 -0
  84. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_debug.py +0 -0
  85. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_dot.py +0 -0
  86. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_einsum.py +0 -0
  87. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_fp8.py +0 -0
  88. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_hof.py +0 -0
  89. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_nn.py +0 -0
  90. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_ops.py +0 -0
  91. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_parsing.py +0 -0
  92. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_partitioning.py +0 -0
  93. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_pool.py +0 -0
  94. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_random.py +0 -0
  95. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_rearrange.py +0 -0
  96. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_scan.py +0 -0
  97. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_specialized_fns.py +0 -0
  98. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_tree_util.py +0 -0
  99. {haliax-1.4.dev299 → haliax-1.4.dev300}/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  name: Pre-Commit
2
2
 
3
- on: [push]
3
+ on: [push, pull_request]
4
4
 
5
5
  jobs:
6
6
  build:
@@ -0,0 +1,36 @@
1
+ name: Run Levanter Tests
2
+
3
+ on: [pull_request]
4
+
5
+ jobs:
6
+ build:
7
+
8
+ runs-on: ubuntu-latest
9
+
10
+ steps:
11
+ - uses: actions/checkout@v3
12
+ - name: Set up Python 3.10.11
13
+ uses: actions/setup-python@v4
14
+ with:
15
+ python-version: 3.10.11
16
+ - name: Install dependencies
17
+ run: |
18
+ python -m pip install --upgrade pip
19
+ pip install flake8 pytest
20
+ pip install --upgrade "jax[cpu]==0.4.26" "jaxlib[cpu]==0.4.26"
21
+
22
+ - name: Install Levanter from source
23
+ run: |
24
+ cd ..
25
+ git clone https://github.com/stanford-crfm/levanter.git
26
+ cd levanter
27
+ pip install -e .
28
+ - name: Install Haliax on top
29
+ run: |
30
+ # install second since levanter will install a built version of haliax
31
+ cd ../haliax
32
+ pip install .[dev]
33
+ - name: Test levanter with pytest
34
+ run: |
35
+ cd ../levanter
36
+ XLA_FLAGS=--xla_force_host_platform_device_count=8 PYTHONPATH=tests:src:../src pytest tests -m "not entry and not slow"
@@ -1,21 +1,18 @@
1
1
  name: Run Tests
2
2
 
3
- on: [push]
3
+ on: [push, pull_request]
4
4
 
5
5
  jobs:
6
6
  build:
7
7
 
8
8
  runs-on: ubuntu-latest
9
- strategy:
10
- matrix:
11
- python-version: ["3.10.11"]
12
9
 
13
10
  steps:
14
11
  - uses: actions/checkout@v3
15
- - name: Set up Python ${{ matrix.python-version }}
12
+ - name: Set up Python 3.10.11
16
13
  uses: actions/setup-python@v4
17
14
  with:
18
- python-version: ${{ matrix.python-version }}
15
+ python-version: 3.10.11
19
16
  - name: Install dependencies
20
17
  run: |
21
18
  python -m pip install --upgrade pip
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: haliax
3
- Version: 1.4.dev299
3
+ Version: 1.4.dev300
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/
@@ -0,0 +1 @@
1
+ __version__ = "1.4.dev300"
@@ -1 +0,0 @@
1
- __version__ = "1.4.dev299"
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