Code

feat: initial nnx port#3

Code

Merged

ariG23498 merged 34 commits intomainml-gde/jflux:mainfrom aritra/modulesml-gde/jflux:aritra/modulesCopy head branch name to clipboardSep 29, 2024

+4,194-107

Commits

Commits on Sep 2, 2024

• chore: inital commit for porting

  ariG23498

  committed
  89aaaeb

• adding torch code from original repository

  ariG23498

  committed
  e5c0c2d

Commits on Sep 3, 2024

• dot pdt attn in jnp

  ariG23498

  committed
  14398a3

• more changes to init

  ariG23498

  committed
  1a00d0c

• start porting

  ariG23498

  committed
  3dfd926

Commits on Sep 4, 2024

• adding everything to uv

  ariG23498

  committed
  3bcd923

• adding ruff and formating

  ariG23498

  committed
  3e09e36

Commits on Sep 8, 2024

• fix(modules/autoencoder): correct invocation of GroupNorm and Conv

  SauravMaheshkar

  committed
  92e2c5b

• fix(modules/layers): correct invocation of nnx modules

  SauravMaheshkar

  committed
  d894b3a

Commits on Sep 13, 2024

• fix: jaxify codebase

  SauravMaheshkar

  committed
  0414a02

• chore: allow for non GPU installation

  SauravMaheshkar

  committed
  38e97b2

• feat(ci): use uv in CI

  SauravMaheshkar

  committed
  363f74e

• fix: tmp mark test as xfail

  SauravMaheshkar

  committed
  4a3948e

• fix: use Sequential for the middle blocks

  SauravMaheshkar

  committed
  29ad2a0

• chore: drop devcontainer

  SauravMaheshkar

  committed
  b4b3c37

• feat: add loop to cli

  SauravMaheshkar

  committed
  8c6a590

• chore: fix just cmd

  SauravMaheshkar

  committed
  393429a

• style(mypy): disable no-redef

  SauravMaheshkar

  committed
  76c05b2

• fix: return numpy tensors from tokenizer

  SauravMaheshkar

  committed
  25c2bf1

• feat: use Array from chex

  SauravMaheshkar

  committed
  52464ef

• docs: update docstrings + use chex

  SauravMaheshkar

  committed
  eadbc6d

• fix: nnx modules use __call__

  SauravMaheshkar

  committed
  cb9a5a7

• feat: jaxify prepare fn

  SauravMaheshkar

  committed
  1bfa089

• fix: nnx modules use __call__

  SauravMaheshkar

  committed
  919aace

• docs: add docstrings to denoise fn

  SauravMaheshkar

  committed
  6c97c8d

• feat: to_device >> device_put

  SauravMaheshkar

  committed
  df0e52b

• feat: add dtypes and param dtypes

  SauravMaheshkar

  committed
  09d1e08

• docs: docstrings for Identity module

  SauravMaheshkar

  committed
  52776d6

• feat: explicitly specify the dtypes for QKNorm in SelfAttention module

  SauravMaheshkar

  committed
  8a64274

Commits on Sep 15, 2024

• feat: add tests for embedding layer

  SauravMaheshkar

  committed
  3d9dc3f

Commits on Sep 16, 2024

• feat: use official flux as optional deps

  SauravMaheshkar

  committed
  b48c3cc

• style: enforce isort

  SauravMaheshkar

  committed
  7775066

• feat: add tests for layers

  SauravMaheshkar

  committed
  af24690

• feat: add tests for modulation and self-attn

  SauravMaheshkar

  committed
  0c59aa1