make StatefulShardedModel callable, add some tests for mlx sharded inference

This commit is contained in:
Alex Cheema
2024-07-13 15:41:15 -07:00
parent 6ee0547eff
commit 850b72d3ea
8 changed files with 103 additions and 2 deletions

1
.gitignore vendored
View File

@@ -1,2 +1,3 @@
__pycache__/
.venv
test_weights.npz