haliax 1.4.dev405__tar.gz → 1.4.dev406__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 (121) hide show
  1. haliax-1.4.dev406/.agents/projects/api_parity.md +215 -0
  2. {haliax-1.4.dev405 → haliax-1.4.dev406}/PKG-INFO +1 -1
  3. haliax-1.4.dev406/src/haliax/__about__.py +1 -0
  4. haliax-1.4.dev405/src/haliax/__about__.py +0 -1
  5. {haliax-1.4.dev405 → haliax-1.4.dev406}/.coveragerc +0 -0
  6. {haliax-1.4.dev405 → haliax-1.4.dev406}/.flake8 +0 -0
  7. {haliax-1.4.dev405 → haliax-1.4.dev406}/.github/workflows/publish_dev.yaml +0 -0
  8. {haliax-1.4.dev405 → haliax-1.4.dev406}/.github/workflows/run_pre_commit.yaml +0 -0
  9. {haliax-1.4.dev405 → haliax-1.4.dev406}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
  10. {haliax-1.4.dev405 → haliax-1.4.dev406}/.github/workflows/run_tests.yaml +0 -0
  11. {haliax-1.4.dev405 → haliax-1.4.dev406}/.gitignore +0 -0
  12. {haliax-1.4.dev405 → haliax-1.4.dev406}/.playbooks/add-types.md +0 -0
  13. {haliax-1.4.dev405 → haliax-1.4.dev406}/.playbooks/wrap-non-named.md +0 -0
  14. {haliax-1.4.dev405 → haliax-1.4.dev406}/.pre-commit-config.yaml +0 -0
  15. {haliax-1.4.dev405 → haliax-1.4.dev406}/.readthedocs.yaml +0 -0
  16. {haliax-1.4.dev405 → haliax-1.4.dev406}/AGENTS.md +0 -0
  17. {haliax-1.4.dev405 → haliax-1.4.dev406}/CONTRIBUTING.md +0 -0
  18. {haliax-1.4.dev405 → haliax-1.4.dev406}/LICENSE +0 -0
  19. {haliax-1.4.dev405 → haliax-1.4.dev406}/README.md +0 -0
  20. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/api.md +0 -0
  21. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/broadcasting.md +0 -0
  22. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/cheatsheet.md +0 -0
  23. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/css/material.css +0 -0
  24. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/css/mkdocstrings.css +0 -0
  25. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/faq.md +0 -0
  26. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/data_parallel_mesh.png +0 -0
  27. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/data_parallel_mesh_replicated.png +0 -0
  28. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_1d.png +0 -0
  29. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_1d_zero.png +0 -0
  30. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_2d.png +0 -0
  31. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
  32. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
  33. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
  34. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
  35. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/figures/device_mesh_2d_zero.png +0 -0
  36. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/fp8.md +0 -0
  37. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/index.md +0 -0
  38. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/indexing.md +0 -0
  39. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/matmul.md +0 -0
  40. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/nn.md +0 -0
  41. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/partitioning.md +0 -0
  42. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/primer.md +0 -0
  43. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/rearrange.ipynb +0 -0
  44. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/rearrange.md +0 -0
  45. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/requirements.txt +0 -0
  46. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/scan.md +0 -0
  47. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/state-dict.md +0 -0
  48. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/tutorial.md +0 -0
  49. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/typing.md +0 -0
  50. {haliax-1.4.dev405 → haliax-1.4.dev406}/docs/vmap.md +0 -0
  51. {haliax-1.4.dev405 → haliax-1.4.dev406}/mkdocs.yml +0 -0
  52. {haliax-1.4.dev405 → haliax-1.4.dev406}/pyproject.toml +0 -0
  53. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/__init__.py +0 -0
  54. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/__init__.py +0 -0
  55. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/compile_utils.py +0 -0
  56. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/dot.py +0 -0
  57. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/einsum.py +0 -0
  58. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/fp8.py +0 -0
  59. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/parsing.py +0 -0
  60. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/rearrange.py +0 -0
  61. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/scan.py +0 -0
  62. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/state_dict.py +0 -0
  63. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/_src/util.py +0 -0
  64. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/axis.py +0 -0
  65. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/core.py +0 -0
  66. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/debug.py +0 -0
  67. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/field.py +0 -0
  68. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/haxtyping.py +0 -0
  69. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/hof.py +0 -0
  70. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/jax_utils.py +0 -0
  71. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/__init__.py +0 -0
  72. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/activations.py +0 -0
  73. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/attention.py +0 -0
  74. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/conv.py +0 -0
  75. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/dropout.py +0 -0
  76. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/embedding.py +0 -0
  77. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/linear.py +0 -0
  78. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/loss.py +0 -0
  79. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/mlp.py +0 -0
  80. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/normalization.py +0 -0
  81. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/pool.py +0 -0
  82. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/nn/scan.py +0 -0
  83. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/ops.py +0 -0
  84. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/partitioning.py +0 -0
  85. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/quantization.py +0 -0
  86. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/random.py +0 -0
  87. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/specialized_fns.py +0 -0
  88. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/state_dict.py +0 -0
  89. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/tree_util.py +0 -0
  90. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/types.py +0 -0
  91. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/util.py +0 -0
  92. {haliax-1.4.dev405 → haliax-1.4.dev406}/src/haliax/wrap.py +0 -0
  93. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/core_test.py +0 -0
  94. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_attention.py +0 -0
  95. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_axis.py +0 -0
  96. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_conv.py +0 -0
  97. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_debug.py +0 -0
  98. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_dot.py +0 -0
  99. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_dtype_typing.py +0 -0
  100. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_einsum.py +0 -0
  101. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_field.py +0 -0
  102. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_fp8.py +0 -0
  103. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_hof.py +0 -0
  104. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_int8.py +0 -0
  105. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_moe_linear.py +0 -0
  106. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_namedarray_typing.py +0 -0
  107. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_nn.py +0 -0
  108. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_ops.py +0 -0
  109. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_parsing.py +0 -0
  110. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_partitioning.py +0 -0
  111. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_pool.py +0 -0
  112. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_random.py +0 -0
  113. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_rearrange.py +0 -0
  114. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_scan.py +0 -0
  115. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_scatter_gather.py +0 -0
  116. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_specialized_fns.py +0 -0
  117. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_state_dict.py +0 -0
  118. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_tree_util.py +0 -0
  119. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_utils.py +0 -0
  120. {haliax-1.4.dev405 → haliax-1.4.dev406}/tests/test_visualize_sharding.py +0 -0
  121. {haliax-1.4.dev405 → haliax-1.4.dev406}/uv.lock +0 -0
@@ -0,0 +1,215 @@
1
+ # JAX NumPy API Parity
2
+ This document tracks JAX NumPy functions not yet wrapped by Haliax.
3
+
4
+ APIs that don't translate well to named tensors are intentionally omitted here. This includes dtype constructors, raw array converters (e.g. `from_dlpack`), indexing helpers like `c_`/`r_`, and functions whose JAX counterparts already work with `NamedArray` out of the box. Basic array construction is handled by `haliax.named`, so functions such as `array` or `asarray` are not listed.
5
+
6
+ ## numpy
7
+ - [ ] `allclose`
8
+ - [ ] `amin`
9
+ - [ ] `append`
10
+ - [ ] `apply_along_axis`
11
+ - [ ] `apply_over_axes`
12
+ - [ ] `argpartition`
13
+ - [ ] `argwhere`
14
+ - [ ] `array_equal`
15
+ - [ ] `array_equiv`
16
+ - [ ] `array_split`
17
+ - [ ] `astype`
18
+ - [ ] `atan2`
19
+ - [ ] `average`
20
+ - [ ] `bartlett`
21
+ - [ ] `bitwise_count`
22
+ - [ ] `bitwise_invert`
23
+ - [ ] `bitwise_left_shift`
24
+ - [ ] `bitwise_right_shift`
25
+ - [ ] `blackman`
26
+ - [ ] `block`
27
+ - [ ] `broadcast_shapes`
28
+ - [ ] `can_cast`
29
+ - [ ] `choose`
30
+ - [ ] `column_stack`
31
+ - [ ] `compress`
32
+ - [ ] `concat`
33
+ - [ ] `convolve`
34
+ - [ ] `copysign`
35
+ - [ ] `corrcoef`
36
+ - [ ] `correlate`
37
+ - [ ] `count_nonzero`
38
+ - [ ] `cov`
39
+ - [ ] `cumulative_prod`
40
+ - [ ] `cumulative_sum`
41
+ - [ ] `delete`
42
+ - [ ] `diag`
43
+ - [ ] `diag_indices`
44
+ - [ ] `diag_indices_from`
45
+ - [ ] `diagflat`
46
+ - [ ] `diagonal`
47
+ - [ ] `diff`
48
+ - [ ] `digitize`
49
+ - [ ] `dsplit`
50
+ - [ ] `dstack`
51
+ - [ ] `ediff1d`
52
+ - [ ] `einsum_path`
53
+ - [ ] `empty`
54
+ - [ ] `empty_like`
55
+ - [ ] `expand_dims`
56
+ - [ ] `extract`
57
+ - [ ] `eye`
58
+ - [ ] `fill_diagonal`
59
+ - [ ] `flatnonzero`
60
+ - [ ] `flip`
61
+ - [ ] `fliplr`
62
+ - [ ] `flipud`
63
+ - [ ] `gcd`
64
+ - [ ] `gradient`
65
+ - [ ] `hamming`
66
+ - [ ] `hanning`
67
+ - [ ] `heaviside`
68
+ - [ ] `histogram`
69
+ - [ ] `histogram2d`
70
+ - [ ] `histogram_bin_edges`
71
+ - [ ] `histogramdd`
72
+ - [ ] `hsplit`
73
+ - [ ] `hstack`
74
+ - [ ] `identity`
75
+ - [ ] `indices`
76
+ - [ ] `inner`
77
+ - [ ] `insert`
78
+ - [ ] `interp`
79
+ - [ ] `intersect1d`
80
+ - [ ] `iscomplexobj`
81
+ - [ ] `isin`
82
+ - [ ] `iterable`
83
+ - [ ] `ix_`
84
+ - [ ] `kaiser`
85
+ - [ ] `kron`
86
+ - [ ] `lcm`
87
+ - [ ] `ldexp`
88
+ - [ ] `lexsort`
89
+ - [ ] `mask_indices`
90
+ - [ ] `median`
91
+ - [ ] `meshgrid`
92
+ - [ ] `mgrid`
93
+ - [ ] `modf`
94
+ - [ ] `moveaxis`
95
+ - [ ] `nan_to_num`
96
+ - [ ] `nanargmax`
97
+ - [ ] `nanargmin`
98
+ - [ ] `nancumprod`
99
+ - [ ] `nancumsum`
100
+ - [ ] `nanmax`
101
+ - [ ] `nanmean`
102
+ - [ ] `nanmedian`
103
+ - [ ] `nanmin`
104
+ - [ ] `nanpercentile`
105
+ - [ ] `nanprod`
106
+ - [ ] `nanquantile`
107
+ - [ ] `nanstd`
108
+ - [ ] `nansum`
109
+ - [ ] `nanvar`
110
+ - [ ] `nonzero`
111
+ - [ ] `ogrid`
112
+ - [ ] `packbits`
113
+ - [ ] `partition`
114
+ - [ ] `percentile`
115
+ - [ ] `permute_dims`
116
+ - [ ] `piecewise`
117
+ - [ ] `place`
118
+ - [ ] `poly`
119
+ - [ ] `polyadd`
120
+ - [ ] `polyder`
121
+ - [ ] `polydiv`
122
+ - [ ] `polyfit`
123
+ - [ ] `polyint`
124
+ - [ ] `polymul`
125
+ - [ ] `polysub`
126
+ - [ ] `polyval`
127
+ - [ ] `pow`
128
+ - [ ] `promote_types`
129
+ - [ ] `put`
130
+ - [ ] `put_along_axis`
131
+ - [ ] `quantile`
132
+ - [ ] `ravel_multi_index`
133
+ - [ ] `reshape`
134
+ - [ ] `resize`
135
+ - [ ] `result_type`
136
+ - [ ] `rollaxis`
137
+ - [ ] `roots`
138
+ - [ ] `rot90`
139
+ - [ ] `select`
140
+ - [ ] `setdiff1d`
141
+ - [ ] `setxor1d`
142
+ - [ ] `spacing`
143
+ - [ ] `squeeze`
144
+ - [ ] `swapaxes`
145
+ - [ ] `take_along_axis`
146
+ - [ ] `transpose`
147
+ - [ ] `trapezoid`
148
+ - [ ] `tri`
149
+ - [ ] `tril_indices`
150
+ - [ ] `tril_indices_from`
151
+ - [ ] `trim_zeros`
152
+ - [ ] `triu_indices`
153
+ - [ ] `triu_indices_from`
154
+ - [ ] `union1d`
155
+ - [ ] `unpackbits`
156
+ - [ ] `unravel_index`
157
+ - [ ] `unstack`
158
+ - [ ] `unwrap`
159
+ - [ ] `vander`
160
+ - [ ] `vsplit`
161
+ - [ ] `vstack`
162
+
163
+ ## fft
164
+ - [ ] `fft`
165
+ - [ ] `fft2`
166
+ - [ ] `fftfreq`
167
+ - [ ] `fftn`
168
+ - [ ] `fftshift`
169
+ - [ ] `hfft`
170
+ - [ ] `ifft`
171
+ - [ ] `ifft2`
172
+ - [ ] `ifftn`
173
+ - [ ] `ifftshift`
174
+ - [ ] `ihfft`
175
+ - [ ] `irfft`
176
+ - [ ] `irfft2`
177
+ - [ ] `irfftn`
178
+ - [ ] `rfft`
179
+ - [ ] `rfft2`
180
+ - [ ] `rfftfreq`
181
+ - [ ] `rfftn`
182
+
183
+ ## linalg
184
+ - [ ] `cholesky`
185
+ - [ ] `cond`
186
+ - [ ] `cross`
187
+ - [ ] `det`
188
+ - [ ] `diagonal`
189
+ - [ ] `eig`
190
+ - [ ] `eigh`
191
+ - [ ] `eigvals`
192
+ - [ ] `eigvalsh`
193
+ - [ ] `inv`
194
+ - [ ] `lstsq`
195
+ - [ ] `matmul`
196
+ - [ ] `matrix_norm`
197
+ - [ ] `matrix_power`
198
+ - [ ] `matrix_rank`
199
+ - [ ] `matrix_transpose`
200
+ - [ ] `norm`
201
+ - [ ] `outer`
202
+ - [ ] `pinv`
203
+ - [ ] `qr`
204
+ - [ ] `slogdet`
205
+ - [ ] `solve`
206
+ - [ ] `svd`
207
+ - [ ] `svdvals`
208
+ - [ ] `tensordot`
209
+ - [ ] `tensorinv`
210
+ - [ ] `tensorsolve`
211
+ - [ ] `trace`
212
+ - [ ] `vdot`
213
+ - [ ] `vecdot`
214
+ - [ ] `vecmat`
215
+ - [ ] `vector_norm`
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: haliax
3
- Version: 1.4.dev405
3
+ Version: 1.4.dev406
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.dev406"
@@ -1 +0,0 @@
1
- __version__ = "1.4.dev405"
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
File without changes
File without changes
File without changes
File without changes