class_docstr: A Keras layer built from a Trax layer.
  
  This subclass of `tf.keras.layers.Layer` takes in a Trax layer as a
  constructor argument and wraps it to be a Keras layer. It uses
  `tf.Variable` to store weights and state (initialized according to the Trax
  layer), and uses the Trax layer's forward function as its forward function.
  
  Consider this code snippet::
  
      keras_layer = AsKeras(trax_layer, initializer_rng=initializer_rng,
                                   rng=rng, rng_updater=rng_updater)
      keras_layer.build(...)  # optional
      outputs = keras_layer(inputs)
  
  (Note that in Keras calling `Layer.build` is optional. If omitted, it will be
  called automatically by `Layer.__call__`.)
  
  If `trax_layer` already has weights at `build` time, the snippet is roughly
  equivalent to::
  
      weights = trax_layer.weights
      state = trax_layer.state
      keras_layer = tf.keras.layers.Layer()
      keras_layer._weights = tf.Variable(weights)
      keras_layer._state = tf.Variable(state)
      keras_layer._rng = tf.Variable(rng)
      outputs, new_state = trax_layer(inputs, keras_layer._weights,
                                      keras_layer._state, keras_layer._rng)
      keras_layer._state.assign(new_state)
      keras_layer._rng.assign(rng_updater(rng))
  
  If `trax_layer` doesn't have weights at `build` time, the snippet is roughly
  equivalent to::
  
      weights, state = trax_layer.init(..., rng=initializer_rng)
      keras_layer = ...
      ...
  
  `AsKeras` uses `tf.Variable` to store weights, not shared with the
  original Trax layer (which uses tensors to store weights), so using
  `AsKeras` may double the memory footprint. This problem can be solved
  by making sure that the Trax layer's weights/state are cleared whenever
  `tf.Variable.assign` (and `tf.Variable.assign_add` etc.) is called, because
  `tf.Variable` is copy-on-write by default.
  
  Mutations in those `tf.Variable`s won't affect the Trax layer's weights, but
  `AsKeras`'s forward function calls the Trax layer's forward function,
  which caches the weights in the Trax layer object, so a forward pass may
  change the weights cached in the original Trax layer.
  
  Note that this class is not thread-safe. If the same `AsKeras` object
  is used in multiple threads, the `tf.Variable` updates may happen in a
  non-deterministic order.
code: def call(self, inputs):
      with math_lib.use_backend(math_lib.Backend.TFNP):
          inputs = math_lib.nested_map(
              functools.partial(
                  _replace_none_batch, batch_size=self._batch_size
              ),
              inputs,
          )
          weights, state, rng = read_values(
              [self._weights, self._state, self._rng]
          )
          inputs, weights, state, rng = to_arrays(
              [inputs, weights, state, rng]
          )
          outputs, new_state = self._trax_layer.pure_fn(
              inputs, weights=weights, state=state, rng=rng
          )
          tf.nest.map_structure(
              lambda v, t: v.assign(t), self._state, new_state
          )
          self._rng.assign(self._rng_updater(rng))
          outputs = to_tensors(outputs)
          return outputs
  
file_path: trax/trax/trax2keras.py
imports: import functools
  import tensorflow.compat.v2
  from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers import base
label_desc: Use image node_id 3 for calling the AsKeras obj's underlying member method code with example usage: obj.call(inputs) and returns: outputs