haliax 1.4.dev317__tar.gz → 1.4.dev318__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.dev317 → haliax-1.4.dev318}/PKG-INFO +4 -1
  2. {haliax-1.4.dev317 → haliax-1.4.dev318}/README.md +3 -0
  3. haliax-1.4.dev318/src/haliax/__about__.py +1 -0
  4. haliax-1.4.dev317/src/haliax/__about__.py +0 -1
  5. {haliax-1.4.dev317 → haliax-1.4.dev318}/.coveragerc +0 -0
  6. {haliax-1.4.dev317 → haliax-1.4.dev318}/.flake8 +0 -0
  7. {haliax-1.4.dev317 → haliax-1.4.dev318}/.github/workflows/publish_dev.yaml +0 -0
  8. {haliax-1.4.dev317 → haliax-1.4.dev318}/.github/workflows/run_pre_commit.yaml +0 -0
  9. {haliax-1.4.dev317 → haliax-1.4.dev318}/.github/workflows/run_quick_levanter_tests.yaml +0 -0
  10. {haliax-1.4.dev317 → haliax-1.4.dev318}/.github/workflows/run_tests.yaml +0 -0
  11. {haliax-1.4.dev317 → haliax-1.4.dev318}/.gitignore +0 -0
  12. {haliax-1.4.dev317 → haliax-1.4.dev318}/.pre-commit-config.yaml +0 -0
  13. {haliax-1.4.dev317 → haliax-1.4.dev318}/.readthedocs.yaml +0 -0
  14. {haliax-1.4.dev317 → haliax-1.4.dev318}/CONTRIBUTING.md +0 -0
  15. {haliax-1.4.dev317 → haliax-1.4.dev318}/LICENSE +0 -0
  16. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/api.md +0 -0
  17. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/broadcasting.md +0 -0
  18. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/cheatsheet.md +0 -0
  19. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/css/material.css +0 -0
  20. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/css/mkdocstrings.css +0 -0
  21. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/faq.md +0 -0
  22. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/data_parallel_mesh.png +0 -0
  23. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/data_parallel_mesh_replicated.png +0 -0
  24. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_1d.png +0 -0
  25. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_1d_zero.png +0 -0
  26. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_2d.png +0 -0
  27. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_2d_batch_partitioned.png +0 -0
  28. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_2d_data_replicated.png +0 -0
  29. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_2d_data_replicated_mlp_partitioned.png +0 -0
  30. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_2d_intermediate_fully_partitioned.png +0 -0
  31. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/figures/device_mesh_2d_zero.png +0 -0
  32. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/fp8.md +0 -0
  33. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/hof.md +0 -0
  34. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/index.md +0 -0
  35. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/indexing.md +0 -0
  36. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/matmul.md +0 -0
  37. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/nn.md +0 -0
  38. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/partitioning.md +0 -0
  39. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/rearrange.ipynb +0 -0
  40. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/rearrange.md +0 -0
  41. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/requirements.txt +0 -0
  42. {haliax-1.4.dev317 → haliax-1.4.dev318}/docs/tutorial.md +0 -0
  43. {haliax-1.4.dev317 → haliax-1.4.dev318}/mkdocs.yml +0 -0
  44. {haliax-1.4.dev317 → haliax-1.4.dev318}/pyproject.toml +0 -0
  45. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/__init__.py +0 -0
  46. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/__init__.py +0 -0
  47. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/compile_utils.py +0 -0
  48. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/dot.py +0 -0
  49. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/einsum.py +0 -0
  50. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/fp8.py +0 -0
  51. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/parsing.py +0 -0
  52. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/rearrange.py +0 -0
  53. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/_src/util.py +0 -0
  54. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/axis.py +0 -0
  55. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/core.py +0 -0
  56. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/debug.py +0 -0
  57. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/hof.py +0 -0
  58. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/jax_utils.py +0 -0
  59. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/__init__.py +0 -0
  60. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/activations.py +0 -0
  61. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/attention.py +0 -0
  62. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/conv.py +0 -0
  63. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/dropout.py +0 -0
  64. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/embedding.py +0 -0
  65. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/linear.py +0 -0
  66. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/loss.py +0 -0
  67. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/mlp.py +0 -0
  68. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/normalization.py +0 -0
  69. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/pool.py +0 -0
  70. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/nn/scan.py +0 -0
  71. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/ops.py +0 -0
  72. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/partitioning.py +0 -0
  73. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/quantization.py +0 -0
  74. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/random.py +0 -0
  75. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/specialized_fns.py +0 -0
  76. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/tree_util.py +0 -0
  77. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/types.py +0 -0
  78. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/util.py +0 -0
  79. {haliax-1.4.dev317 → haliax-1.4.dev318}/src/haliax/wrap.py +0 -0
  80. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/core_test.py +0 -0
  81. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_attention.py +0 -0
  82. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_axis.py +0 -0
  83. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_conv.py +0 -0
  84. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_debug.py +0 -0
  85. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_dot.py +0 -0
  86. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_einsum.py +0 -0
  87. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_fp8.py +0 -0
  88. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_hof.py +0 -0
  89. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_nn.py +0 -0
  90. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_ops.py +0 -0
  91. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_parsing.py +0 -0
  92. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_partitioning.py +0 -0
  93. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_pool.py +0 -0
  94. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_random.py +0 -0
  95. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_rearrange.py +0 -0
  96. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_scan.py +0 -0
  97. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_specialized_fns.py +0 -0
  98. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_tree_util.py +0 -0
  99. {haliax-1.4.dev317 → haliax-1.4.dev318}/tests/test_utils.py +0 -0
{haliax-1.4.dev317 → haliax-1.4.dev318}/PKG-INFO RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: haliax
3
- Version: 1.4.dev317
3
+ Version: 1.4.dev318
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/
@@ -80,6 +80,9 @@ Head = hax.Axis("head", 8) # number of attention heads
80
80
  Key = hax.Axis("key", 64) # key size
81
81
  Embed = hax.Axis("embed", 512) # embedding size
82
82
 
83
+ # alternatively:
84
+ #Pos, KPos, Head, Key, Embed = hax.make_axes(pos=1024, key_pos=1024, head=8, key=64, embed=512)
85
+
83
86
 
84
87
  def attention_scores(Key, KPos, query, key, mask):
85
88
  # how similar is each query to each key
{haliax-1.4.dev317 → haliax-1.4.dev318}/README.md RENAMED
@@ -47,6 +47,9 @@ Head = hax.Axis("head", 8) # number of attention heads
47
47
  Key = hax.Axis("key", 64) # key size
48
48
  Embed = hax.Axis("embed", 512) # embedding size
49
49
 
50
+ # alternatively:
51
+ #Pos, KPos, Head, Key, Embed = hax.make_axes(pos=1024, key_pos=1024, head=8, key=64, embed=512)
52
+
50
53
 
51
54
  def attention_scores(Key, KPos, query, key, mask):
52
55
  # how similar is each query to each key
haliax-1.4.dev318/src/haliax/__about__.py ADDED
@@ -0,0 +1 @@
1
+ __version__ = "1.4.dev318"
haliax-1.4.dev317/src/haliax/__about__.py DELETED
@@ -1 +0,0 @@
1
- __version__ = "1.4.dev317"
{haliax-1.4.dev317 → haliax-1.4.dev318}/.coveragerc RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/.flake8 RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/.gitignore RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/CONTRIBUTING.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/LICENSE RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/docs/api.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/docs/faq.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/docs/fp8.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/docs/hof.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/docs/index.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/docs/matmul.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/docs/nn.md RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/mkdocs.yml RENAMED
File without changes
{haliax-1.4.dev317 → haliax-1.4.dev318}/pyproject.toml RENAMED
File without changes