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 build(self, input_shape):
      with math_lib.use_backend(math_lib.Backend.TFNP):
          # Using `is` instead of `==` following Trax's practice
          if self._trax_layer.weights is base.EMPTY_WEIGHTS:
              sanitized_input_shape = math_lib.nested_map(
                  functools.partial(
                      _replace_none_batch, batch_size=self._batch_size
                  ),
                  input_shape,
              )
              weights, state = self._trax_layer.init(
                  tensor_shapes_to_shape_dtypes(
                      sanitized_input_shape, self.dtype
                  ),
                  rng=self._initializer_rng,
              )
          else:
              weights = self._trax_layer.weights
              state = self._trax_layer.state
          # Note: `weights` may contain `EMPTY_WEIGHTS`
          self._weights = math_lib.nested_map(
              functools.partial(tf.Variable, trainable=True), weights
          )
          self._state = math_lib.nested_map(
              functools.partial(tf.Variable, trainable=False), state
          )
          self._rng = tf.Variable(
              self._forward_rng_init, trainable=False
          )
      super().build(input_shape)
  
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 2 for calling the AsKeras obj's underlying member method code with example usage: obj.build(input_shape) without return types