Code

(closes #35) Changes for jax-metal#52

Code

Open

nix wants to merge 4 commits intoxjdr-alt:mainxjdr-alt/entropix:mainfrom nix:mainnix/entropix:mainCopy head branch name to clipboard

+32-11

Commits

Commits on Oct 8, 2024

• Convert RoPE to float since jax-metal doesn't support complex.

  nix

  committed
  fab8bde

  Copy full SHA for fab8bde

• Use jax-metal device if available.

  nix

  committed
  35d9a83

  Copy full SHA for 35d9a83

• Reimplement top_k for jax-metal to avoid nonworking jax.lax.top_k.

  Show description for 5bbaa5b
  nix

  committed
  5bbaa5b

  Copy full SHA for 5bbaa5b

• Flip the sense of the enable flag.

  nix

  committed
  5539d78

  Copy full SHA for 5539d78