trax
trax copied to clipboard
google
Trax — Deep Learning with Clear Code and Speed
### Description I was trying to perform a simple optimization using `Parallel` combinator other than `Serial`. By using Dense layers in Serial, it is trivial to optimize. However, I could...
How can I get access to current metrics in a training loop? I want to do custom logging while the model runs - specifically I want to log to Weights...
### Description When providing inputs with a constant shape - for instance imagenet32 where examples are always of length *3072*, but it also applies e.g. to this config: https://github.com/google/trax/blob/master/trax/supervised/configs/reformer_imagenet64.gin and...
### Description Importing trax takes too long. ### Environment information ``` OS: ubuntu 18.04 $ pip freeze | grep trax trax==1.3.7 $ pip freeze | grep tensor mesh-tensorflow==0.1.18 tensorboard==2.4.0 tensorboard-plugin-wit==1.7.0...
### Description I have gpu, but trax just can not find it, it report that: WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.) ......
### Description ... ### Environment information ``` OS: $ pip freeze | grep trax trax==1.3.6 $ pip freeze | grep tensor mesh-tensorflow==0.1.16 tensor2tensor==1.15.7 tensorboard==2.3.0 tensorboard-plugin-wit==1.7.0 tensorflow==2.3.0 tensorflow-addons==0.11.2 tensorflow-datasets==3.2.1 tensorflow-estimator==2.3.0 tensorflow-gan==2.0.0...
### Description a tracer error occurs inside jax while initializing pretrained weights in model.init_from_file() function inside this notebook https://colab.research.google.com/github/google/trax/blob/master/trax/models/reformer/text_generation.ipynb --------------------------------------------------------------------------- LayerError Traceback (most recent call last) in () 2 model...
### Description The crossentrophy loss in Trax 1.3.7 gave strange reading. Is this normal?  The same code was running in Trax 1.3.6 with this result.  The loss layer...
### Description My network has m inputs (X1, X2, ...Xm), and n prediction heads (Y1, Y2, ...Yn). Each of the input and output may have different shapes. Each of the...
### Description https://github.com/google/trax/blob/master/trax/models/reformer/text_generation.ipynb runs out of memory on a v3-8 tpu ... ### Environment information ``` OS: $ pip freeze | grep trax trax==1.3.7 $ pip freeze | grep tensor...
