<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2020 Google LLC.
Step1:  Reformer
Step2:  Load example data and model
Step3:  Sample from the model
Now we're ready to sample from the pre-trained Reformer model. Unlike during training, sampling processes the images one pixel and channel value at a time. The TPU colab runtime has 8 cores so we can sample 8 images in parallel.
Step4:  Sampling is an inherently serial process and will take up to 9 minutes to run. A good chunk of that time will be spent on JIT-compiling the code, though, so the code cell below will finish faster when re-run for a second time.
<ASSISTANT_TASK:>
Python Code:
# 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.
Explanation: Copyright 2020 Google LLC.
End of explanation
# Install JAX. This custom build raises the TPU timeout threshold, because the
# default limit of 2 minutes is too short for sampling very long sequences.
!gsutil cp gs://trax-ml/reformer/jaxlib-0.1.39-cp36-none-manylinux2010_x86_64.whl .
!gsutil cp gs://trax-ml/reformer/jax-0.1.59-cp36-none-manylinux2010_x86_64.whl .
!pip install --upgrade -q ./jaxlib-0.1.39-cp36-none-manylinux2010_x86_64.whl
!pip install --upgrade -q ./jax-0.1.59-cp36-none-manylinux2010_x86_64.whl
# Make sure the Colab Runtime is set to Accelerator: TPU.
import requests
import os
if 'TPU_DRIVER_MODE' not in globals():
  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20191206'
  resp = requests.post(url)
  TPU_DRIVER_MODE = 1
# The following is required to use TPU Driver as JAX's backend.
from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
print(config.FLAGS.jax_backend_target)
!pip install --upgrade -q gin git+https://github.com/google/trax.git@v1.2.3
from tensorflow.compat.v1.io.gfile import GFile
import gin
import os
import jax
import trax
from trax.models.beam_search import Search
from trax.supervised import inputs
import numpy as np
import jax.numpy as jnp
from scipy.special import softmax
%matplotlib inline
from matplotlib import pyplot as plt
Explanation: Reformer: Image Generation 
This notebook was designed to run on TPU.
To use TPUs in Colab, click "Runtime" on the main menu bar and select Change runtime type. Set "TPU" as the hardware accelerator.
End of explanation
# Normally we train on the full imagenet64 training set, which is quite large so
# we won't be loading it from this notebook. Instead, let's just load a few PNG
# images to use in our data pipeline.
DATA = []
for i in range(8):
  img = plt.imread(GFile('gs://trax-ml/reformer/img{}.png'.format(i), 'rb'))
  # Convert from RGBA floating-point to RGB integer representation.
  img = np.asarray(img[:, :, :3] * 255, dtype=np.int32)
  DATA.append(img)
# We can examine one of the images to make sure we've loaded it correctly.
plt.figure(figsize=(1.5, 1.5))
plt.axis('off')
plt.imshow(DATA[0])
# We'll be using a pre-trained 12-layer Reformer model.
# First, load the config (which sets all needed hyperparameters).
!gsutil cp gs://trax-ml/reformer/imgnet64/config.gin ./config.gin
gin.parse_config_file('./config.gin')
# Now we construct a ReformerLM instance and load the pre-trained weights.
# The 'predict' mode configures the model to accept single tokens at a time,
# instead of feeding in a complete image all at once.
model_infer = trax.models.ReformerLM(mode='predict')
model_infer.init_from_file(
    'gs://trax-ml/reformer/imgnet64/model.pkl', weights_only=True)
Explanation: Load example data and model
End of explanation
sampling_decoder = Search(
    trax.models.ReformerLM,
    model_infer.weights,
    temperature=1.0,
    max_decode_len=32*64*3,
    )
Explanation: Sample from the model
Now we're ready to sample from the pre-trained Reformer model. Unlike during training, sampling processes the images one pixel and channel value at a time. The TPU colab runtime has 8 cores so we can sample 8 images in parallel.
End of explanation
flat_prompt = []
for i, img in enumerate(DATA[:trax.fastmath.device_count()]):
  img = img.reshape((-1, 64, 3))[:32, :, :]
  flat_prompt.append(img.reshape((-1,)))
prompt = np.stack(flat_prompt, 0)
print("Prompt:")
plt.figure(figsize=(10, 10*8))
for i in range(prompt.shape[0]):
  plt.subplot(1, 8, i+1)
  plt.axis('off')
  plt.imshow(prompt[i].reshape((-1, 64, 3)), aspect='equal')
plt.show()
seqs, scores = sampling_decoder.decode(targets_prefix=prompt, batch_size=8)
print("Sampled completions:")
plt.figure(figsize=(10, 10*8))
for i in range(prompt.shape[0]):
  plt.subplot(1, 8, i+1)
  plt.axis('off')
  plt.imshow(seqs[i, -1].reshape((-1, 64, 3)), aspect='equal')
plt.figure(figsize=(10, 10*8))
for i in range(prompt.shape[0]):
  plt.subplot(1, 8, i+1)
  plt.axis('off')
  img = jnp.concatenate([prompt[i], seqs[i, -1]], -1)
  plt.imshow(img.reshape((-1, 64, 3)), aspect='equal')
Explanation: Sampling is an inherently serial process and will take up to 9 minutes to run. A good chunk of that time will be spent on JIT-compiling the code, though, so the code cell below will finish faster when re-run for a second time.
End of explanation
<END_TASK>