gitextract_ho0gv3se/ ├── .github/ │ └── workflows/ │ └── ci.yml ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── WORKSPACE ├── docs/ │ ├── .gitignore │ ├── Makefile │ ├── api.rst │ ├── conf.py │ ├── ext/ │ │ ├── BUILD │ │ ├── link_tf_api.py │ │ └── link_tf_api_test.py │ ├── index.rst │ ├── modules.rst │ ├── references.bib │ └── requirements.txt ├── examples/ │ ├── BUILD │ ├── README.md │ ├── distributed_cifar10.ipynb │ ├── functional_mlp_mnist.py │ ├── little_gan_on_mnist.ipynb │ ├── mlp_on_mnist.ipynb │ ├── simple_mnist.py │ ├── simple_mnist_test.py │ └── vqvae_example.ipynb ├── readthedocs.yml ├── requirements-test.txt ├── requirements-tf.txt ├── requirements.txt ├── setup.py ├── sonnet/ │ ├── BUILD │ ├── __init__.py │ ├── distribute.py │ ├── functional.py │ ├── initializers.py │ ├── mixed_precision.py │ ├── nets/ │ │ ├── BUILD │ │ ├── __init__.py │ │ └── resnet.py │ ├── optimizers.py │ ├── pad.py │ ├── regularizers.py │ └── src/ │ ├── BUILD │ ├── __init__.py │ ├── axis_norm.py │ ├── axis_norm_test.py │ ├── base.py │ ├── base_test.py │ ├── batch_apply.py │ ├── batch_apply_test.py │ ├── batch_norm.py │ ├── batch_norm_test.py │ ├── bias.py │ ├── bias_test.py │ ├── build.py │ ├── build_defs.bzl │ ├── build_test.py │ ├── conformance/ │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── api_test.py │ │ ├── build_test.py │ │ ├── checkpoint_test.py │ │ ├── checkpoints/ │ │ │ ├── BUILD │ │ │ ├── README.md │ │ │ ├── base_batch_norm_1x2x2x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── base_batch_norm_scale_offset_1x2x2x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── batch_norm_1x2x2x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── batch_norm_scale_offset_1x2x2x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── batch_norm_training_1x2x2x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── bias_3x3x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── cifar10_convnet_2x3_2x2_1x3x3x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv1d_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv1d_lstm_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv1d_transpose_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv2d_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv2d_lstm_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv2d_transpose_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv3d_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv3d_lstm_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── conv3d_transpose_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── cross_replica_batch_norm_1x2x2x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── depthwise_conv2d_3x3_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── dropout/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── ema_2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── embed_100_100/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── generate.py │ │ │ ├── group_norm_2_1x3x4/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── gru_1/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── instance_norm_1_1x3_2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── layer_norm_1_1x3_2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── linear_1x1/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── linear_nobias_1x1/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── lstm_1/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── lstm_8_projected_1/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── mean_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00002 │ │ │ │ ├── checkpoint-1.data-00001-of-00002 │ │ │ │ └── checkpoint-1.index │ │ │ ├── mlp_3x4x5_1x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── mlp_nobias_3x4x5_1x3/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── resnet50/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── sum_2x2/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00002 │ │ │ │ ├── checkpoint-1.data-00001-of-00002 │ │ │ │ └── checkpoint-1.index │ │ │ ├── trainable_state/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── unrolled_lstm_1/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── vanilla_rnn_8/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── vqvae/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ ├── vqvae_ema_eval/ │ │ │ │ ├── checkpoint │ │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ │ └── checkpoint-1.index │ │ │ └── vqvae_ema_train/ │ │ │ ├── checkpoint │ │ │ ├── checkpoint-1.data-00000-of-00001 │ │ │ └── checkpoint-1.index │ │ ├── copy_test.py │ │ ├── descriptors.py │ │ ├── descriptors_test.py │ │ ├── distribute_test.py │ │ ├── doctest_test.py │ │ ├── function_test.py │ │ ├── goldens.py │ │ ├── goldens_test.py │ │ ├── keras_test.py │ │ ├── optimizer_test.py │ │ ├── pickle_test.py │ │ ├── saved_model_test.py │ │ ├── tensorflow1_test.py │ │ └── xla_test.py │ ├── conv.py │ ├── conv_test.py │ ├── conv_transpose.py │ ├── conv_transpose_test.py │ ├── custom_getter.py │ ├── custom_getter_test.py │ ├── deferred.py │ ├── deferred_test.py │ ├── depthwise_conv.py │ ├── depthwise_conv_test.py │ ├── distribute/ │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── distributed_batch_norm.py │ │ ├── distributed_batch_norm_test.py │ │ ├── replicator.py │ │ ├── replicator_test.py │ │ └── replicator_test_utils.py │ ├── dropout.py │ ├── dropout_test.py │ ├── embed.py │ ├── embed_test.py │ ├── functional/ │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── haiku.py │ │ ├── haiku_test.py │ │ ├── jax.py │ │ ├── jax_test.py │ │ ├── optimizers.py │ │ ├── optimizers_test.py │ │ └── utils.py │ ├── group_norm.py │ ├── group_norm_test.py │ ├── initializers.py │ ├── initializers_test.py │ ├── leaky_clip_by_value.py │ ├── leaky_clip_by_value_test.py │ ├── linear.py │ ├── linear_test.py │ ├── metrics.py │ ├── metrics_test.py │ ├── mixed_precision.py │ ├── mixed_precision_test.py │ ├── moving_averages.py │ ├── moving_averages_test.py │ ├── nets/ │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── cifar10_convnet.py │ │ ├── cifar10_convnet_test.py │ │ ├── dnc/ │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── control.py │ │ │ ├── control_test.py │ │ │ ├── read.py │ │ │ ├── read_test.py │ │ │ ├── util.py │ │ │ ├── util_test.py │ │ │ ├── write.py │ │ │ └── write_test.py │ │ ├── mlp.py │ │ ├── mlp_test.py │ │ ├── resnet.py │ │ ├── resnet_test.py │ │ ├── vqvae.py │ │ └── vqvae_test.py │ ├── once.py │ ├── once_test.py │ ├── optimizers/ │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── adam.py │ │ ├── adam_test.py │ │ ├── momentum.py │ │ ├── momentum_test.py │ │ ├── optimizer_tests.py │ │ ├── optimizer_utils.py │ │ ├── rmsprop.py │ │ ├── rmsprop_test.py │ │ ├── sgd.py │ │ └── sgd_test.py │ ├── pad.py │ ├── pad_test.py │ ├── parallel_linear.py │ ├── parallel_linear_test.py │ ├── recurrent.py │ ├── recurrent_test.py │ ├── regularizers.py │ ├── regularizers_test.py │ ├── reshape.py │ ├── reshape_test.py │ ├── scale_gradient.py │ ├── scale_gradient_test.py │ ├── sequential.py │ ├── sequential_test.py │ ├── test_utils.py │ ├── types.py │ ├── utils.py │ └── utils_test.py └── test.sh