<jupyter_start><jupyter_text>Using Trax with TensorFlow NumPy and KerasThis notebook ([run it in colab](https://colab.research.google.com/github/google/trax/blob/master/trax/tf_numpy_and_keras.ipynb)) shows how you can run [Trax](https://trax-ml.readthedocs.io/en/latest/) directly with [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy). You will also see how to use Trax layers and models inside [Keras](https://keras.io/) so you can use Trax in production, e.g., with [TensorFlow.js](https://www.tensorflow.org/js/) or [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving).  1. **Trax with TensorFlow NumPy**: use Trax with [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) without any code changes  1. **Convert Trax to Keras**: how to get a [Keras](https://keras.io/) layer for your Trax model and use it  1. **Exporting Trax Models for Deployment**: how to export Trax models to [TensorFlow SavedModel](https://www.tensorflow.org/guide/saved_model)    1. Trax with TensorFlow NumPyIn Trax, all computations rely on accelerated math operations happening in the `fastmath` module. This module can use different backends for acceleration. One of them is [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) which uses [TensorFlow 2](https://www.tensorflow.org/) to accelerate the computations.The backend can be set using a call to `trax.fastmath.set_backend` as you'll see below. Currently available backends are `jax` (default), `tensorflow-numpy` and `numpy` (for debugging). The `tensorflow-numpy` backend uses [TensorFlow Numpy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) for executing `fastmath` functions on TensorFlow, while the `jax` backend calls [JAX](https://github.com/google/jax) which lowers to TensorFlow XLA.You may see that `tensorflow-numpy` and `jax` backends show different speed and memory characteristics. You may also see different error messages when debugging since it might expose you to the internals of the backends. However for the most part, users can choose a backend and not worry about the internal details of these backends.Let's train the sentiment analysis model from the [Trax intro](https://colab.research.google.com/github/google/trax/blob/master/trax/intro.ipynb) using TensorFlow NumPy to see how it works. **General Setup**Execute the following few cells (once) before running any of the code samples.<jupyter_code>#@title
# Copyright 2020 Google LLC.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.



# Install and import Trax
!pip install -q -U git+https://github.com/google/trax@master

import os
import numpy as np
import trax<jupyter_output><jupyter_text>Here is how you can set the fastmath backend to `tensorflow-numpy` and verify that it's been set.<jupyter_code># Use the tensorflow-numpy backend.
trax.fastmath.set_backend('tensorflow-numpy')
print(trax.fastmath.backend_name())
# Create data streams.
train_stream = trax.data.TFDS('imdb_reviews', keys=('text', 'label'), train=True)()
eval_stream = trax.data.TFDS('imdb_reviews', keys=('text', 'label'), train=False)()

data_pipeline = trax.data.Serial(
    trax.data.Tokenize(vocab_file='en_8k.subword', keys=[0]),
    trax.data.Shuffle(),
    trax.data.FilterByLength(max_length=2048, length_keys=[0]),
    trax.data.BucketByLength(boundaries=[  32, 128, 512, 2048],
                             batch_sizes=[512, 128,  32,    8, 1],
                             length_keys=[0]),
    trax.data.AddLossWeights()
  )
train_batches_stream = data_pipeline(train_stream)
eval_batches_stream = data_pipeline(eval_stream)

# Print example shapes.
example_batch = next(train_batches_stream)
print(f'batch shapes = {[x.shape for x in example_batch]}')
# Create the model.
from trax import layers as tl

model = tl.Serial(
    tl.Embedding(vocab_size=8192, d_feature=256),
    tl.Mean(axis=1),  # Average on axis 1 (length of sentence).
    tl.Dense(2),      # Classify 2 classes.
)

# You can print model structure.
print(model)
# Train the model.
from trax.supervised import training

# Training task.
train_task = training.TrainTask(
    labeled_data=train_batches_stream,
    loss_layer=tl.WeightedCategoryCrossEntropy(),
    optimizer=trax.optimizers.Adam(0.01),
    n_steps_per_checkpoint=500,
)

# Evaluaton task.
eval_task = training.EvalTask(
    labeled_data=eval_batches_stream,
    metrics=[tl.WeightedCategoryCrossEntropy(), tl.WeightedCategoryAccuracy()],
    n_eval_batches=20  # For less variance in eval numbers.
)

# Training loop saves checkpoints to output_dir.
output_dir = os.path.expanduser('~/output_dir/')
training_loop = training.Loop(model,
                              train_task,
                              eval_tasks=[eval_task],
                              output_dir=output_dir)

# Run 2000 steps (batches).
training_loop.run(2000)
# Run on an example.
example_input = next(eval_batches_stream)[0][0]
example_input_str = trax.data.detokenize(example_input, vocab_file='en_8k.subword')
print(f'example input_str: {example_input_str}')
sentiment_activations = model(example_input[None, :])  # Add batch dimension.
print(f'Model returned sentiment activations: {np.asarray(sentiment_activations)}')<jupyter_output>example input_str: The movie features another exceptional collaboration between director <NAME> and cinematographer <NAME>, the first after Toland worked on Citizen Kane. But the talent of both these men was focused on achieving a perfectly crafted movie, understood in the good old American sense as a great story. The technical aspects of the movie are covered so as the viewer gets absorbed into the action that takes place on the screen without submitting to the power of the image. Technique is seen as a vehicle of representation unlike in Citizen Kane where Welles' baroque style almost drew the attention from the story to the way the story was told. One of my favorite moves with deep focus in this film is the drama conveyed by the returning home welcoming of Homer and Al. If Homer's girl, Wilma comes towards him perfectly in focus, Al goes over to his wife also perfectly in focus. This is a brilliant move because it shows only through the use of the image the nature of th[...]<jupyter_text>2. Convert Trax to KerasThanks to [TensorFlow NumPy](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) you can convert the model you just trained into a [Keras](https://keras.io/) layer using `trax.AsKeras`. This allows you to:* use Trax layers inside Keras models* run Trax models with existing Keras input pipelines* export Trax models to [TensorFlow SavedModel](https://www.tensorflow.org/guide/saved_model)When creating  a Keras layer from a Trax one, the Keras layer weights will get initialized to the ones the Trax layer had at the moment of creation. In this way, you can create Keras layers from pre-trained Trax models and save them as SavedModel as shown below.<jupyter_code># Convert the model into a Keras layer, use the weights from model.
keras_layer = trax.AsKeras(model)
print(keras_layer)

# Run the Keras layer to verify it returns the same result.
sentiment_activations = keras_layer(example_input[None, :])
print(f'Keras returned sentiment activations: {np.asarray(sentiment_activations)}')
import tensorflow as tf

# Create a full Keras  model using the layer from Trax.
inputs = tf.keras.Input(shape=(None,), dtype='int32')
hidden = keras_layer(inputs) 
# You can add other Keras layers here operating on hidden.
outputs = hidden
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
print(keras_model)

# Run the Keras model to verify it returns the same result.
sentiment_activations = keras_model(example_input[None, :])
print(f'Keras returned sentiment activations: {np.asarray(sentiment_activations)}')<jupyter_output>Keras returned sentiment activations: [[-1.6396211  1.6328843]]<jupyter_text>3. Exporting Trax Models for DeploymentYou can export the Keras model to disk as [TensorFlow SavedModel](https://www.tensorflow.org/guide/saved_model). It's as simple as calling `keras_model.save` and allows you to use models with TF tools [TensorFlow.js](https://www.tensorflow.org/js/), [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving) and [TensorFlow Lite](https://www.tensorflow.org/lite).<jupyter_code># Save the Keras model to output_dir.
model_file = os.path.join(output_dir, "model_checkpoint")
keras_model.save(model_file)

# Load the model from SavedModel.
loaded_model = tf.keras.models.load_model(model_file)

# Run the loaded model to verify it returns the same result.
sentiment_activations = loaded_model(example_input[None, :])
print(f'Keras returned sentiment activations: {np.asarray(sentiment_activations)}')<jupyter_output>Keras returned sentiment activations: [[-1.6396211  1.6328843]]
{'fraction_non_alphanumeric': 0.09481560662747195, 'fraction_numerical': 0.011651523249599146, 'mean_word_length': 4.887979861548144, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 23, 'lorem ipsum': 0, 'www.': 15, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '447689', 'license': 'Apache-2.0', 'n_tokens_mistral': 2864, 'n_tokens_neox': 2745, 'n_words': 964, 'path': 'trax/tf_numpy_and_keras.ipynb', 'repo_name': 'thoo/trax', 'size': '9368'}