google/flax

Pull Requests Commits

Simplify the type check

kaixih committed 2 years ago

3a0e87ab

fixed mypy errors at HEAD

chiamp committed 2 years ago

e7e33e1e

Update the algorithm to compute fp8 scales

kaixih committed 2 years ago

f4cb269e

Merge pull request #3435 from kaixih:update_fp8_support

a-googler committed 2 years ago

738078c6

Updated docstring for `axis_name` arg to say it's only used for `pmap`

chiamp committed 2 years ago

5557649f

updated lm1b example with jit

chiamp committed 2 years ago

558be01e

Update the fp8 support

kaixih committed 2 years ago

d24db613

Make a number of functions visible.

a-googler committed 2 years ago

fd8fd76a

Merge pull request #3429 from levskaya:abstractrng

a-googler committed 2 years ago

abe205be

Merge pull request #3432 from chiamp:torch

a-googler committed 2 years ago

bf2928c8

Supported computing mean and variance of BatchNorm using mask.

a-googler committed 2 years ago

41273d18

updated torch ConvTranspose

chiamp committed 2 years ago

e3633935

updated torch ConvTranspose

chiamp committed 2 years ago

6c208fc6

Simplify abstract rng creation in param shape-check.

levskaya committed 2 years ago

578f7d44

fully deprecated old RNN api

chiamp committed 2 years ago

8a86bf40

Merge pull request #3431 from google:fix-tokenizers

a-googler committed 2 years ago

56da5e46

remove transformers dependency

cgarciae committed 2 years ago

39ad403e

flax: avoid deprecated jax.random.default_prng_impl()

Jake VanderPlas committed 2 years ago

abab11ff

Merge pull request #3427 from levskaya:typefix

a-googler committed 2 years ago

b0999873

Ignore transient chex deprecationwarning.

levskaya committed 2 years ago

47d1c514

Replaces pjit with jit in spmd.py

marcvanzee committed 2 years ago

a8bc3673

Add instrumentation for Flax checkpoint save and restore calls.

a-googler committed 2 years ago

31eedf16

Merge pull request #3419 from chiamp:gru

a-googler committed 2 years ago

108a58ee

Merge pull request #3410 from chiamp:ruff

a-googler committed 2 years ago

f36dccce

fixed GRU docstring

chiamp committed 2 years ago

766485a5

Merge pull request #3415 from google:fix-rdt

a-googler committed 2 years ago

71429fec

remove pdf target

cgarciae committed 2 years ago

afa7e337

Merge pull request #3403 from chiamp:fix_spectralnorm

a-googler committed 2 years ago

b39e70ab

Merge pull request #3384 from chiamp:dropout

a-googler committed 2 years ago

cc1d4eb3

Make the injectable `dot_general` optional and not pre-initialized so that method interceptor works.

IvyZX committed 2 years ago

fac0cef3

Newer Older