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.
docstr: Creates a Keras layer wrapping around a Trax layer.
  
  Args:
    trax_layer: an object of class `trax.layers.Layer`, the trax layer to
      wrap.
    batch_size: (optional) an integer, the batch size that this Keras layer
      will be used on. Keras sometimes needs to generate a TF graph for a
      layer (e.g. for acceleration or checkpointing). The inputs used to trace
      the graph will have `None` as the length of their batch dimensions, so
      as to generate a graph that can handle any batch size. Some Trax layers
      can't handle tensors whose shapes contain `None`. If `batch_size` is set
      to an integer, the graph will be traced with `batch_size` as the batch
      size instead of `None`. Note that in this case the graph (and the Keras
      layer) can only be used on a specific batch size. If you want to use a
      different batch size, you need to create another `AsKeras` object
      with a different `batch_size`.
    initializer_rng: (optional) an RNG key used to create the weights and
      state if `trax_layer` doesn't have them. If `None`,
      `trax.fastmath.random.get_prng(0)` will be used.
    rng: (optional) an RNG key for the forward function (aka the "forward
      key"). If `None`, `trax.fastmath.random.get_prng(0)` will be used.
    rng_updater: (optional) a function of type rng_key -> rng_key, used to
      update the forward key after each forward pass. If `None`, the function
      `lambda x: trax.fastmath.random.split(x, 1)[0]` will be used, which
      advances the RNG key.
    dtype: (optional) the dtype of the inputs. See the `dtype` argument of
      `tf.keras.layers.Layer.__init__` for details.
docstr_tok: ["Creates","a","Keras","layer","wrapping","around","a","Trax","layer",".","Args",":","trax_layer",":","an","object","of","class","`","trax.layers.Layer","`",",","the","trax","layer","to","wrap",".","batch_size",":","(","optional",")","an","integer",",","the","batch","size","that","this","Keras","layer","will","be","used","on",".","Keras","sometimes","needs","to","generate","a","TF","graph","for","a","layer","(","e.g",".","for","acceleration","or","checkpointing",")",".","The","inputs","used","to","trace","the","graph","will","have","`","None","`","as","the","length","of","their","batch","dimensions",",","so","as","to","generate","a","graph","that","can","handle","any","batch","size",".","Some","Trax","layers","ca","n't","handle","tensors","whose","shapes","contain","`","None","`",".","If","`","batch_size","`","is","set","to","an","integer",",","the","graph","will","be","traced","with","`","batch_size","`","as","the","batch","size","instead","of","`","None","`",".","Note","that","in","this","case","the","graph","(","and","the","Keras","layer",")","can","only","be","used","on","a","specific","batch","size",".","If","you","want","to","use","a","different","batch","size",",","you","need","to","create","another","`","AsKeras","`","object","with","a","different","`","batch_size","`",".","initializer_rng",":","(","optional",")","an","RNG","key","used","to","create","the","weights","and","state","if","`","trax_layer","`","does","n't","have","them",".","If","`","None","`",",","`","trax.fastmath.random.get_prng","(","0",")","`","will","be","used",".","rng",":","(","optional",")","an","RNG","key","for","the","forward","function","(","aka","the","``","forward","key","''",")",".","If","`","None","`",",","`","trax.fastmath.random.get_prng","(","0",")","`","will","be","used",".","rng_updater",":","(","optional",")","a","function","of","type","rng_key","-",">","rng_key",",","used","to","update","the","forward","key","after","each","forward","pass",".","If","`","None","`",",","the","function","`","lambda","x",":","trax.fastmath.random.split","(","x",",","1",")","[","0","]","`","will","be","used",",","which","advances","the","RNG","key",".","dtype",":","(","optional",")","the","dtype","of","the","inputs",".","See","the","`","dtype","`","argument","of","`","tf.keras.layers.Layer.__init__","`","for","details","."]
code: def __init__(
      self,
      trax_layer,
      batch_size=None,
      initializer_rng=None,
      rng=None,
      rng_updater=None,
      dtype=None,
  ):
      """Creates a Keras layer wrapping around a Trax layer.
  
      Args:
        trax_layer: an object of class `trax.layers.Layer`, the trax layer to
          wrap.
        batch_size: (optional) an integer, the batch size that this Keras layer
          will be used on. Keras sometimes needs to generate a TF graph for a
          layer (e.g. for acceleration or checkpointing). The inputs used to trace
          the graph will have `None` as the length of their batch dimensions, so
          as to generate a graph that can handle any batch size. Some Trax layers
          can't handle tensors whose shapes contain `None`. If `batch_size` is set
          to an integer, the graph will be traced with `batch_size` as the batch
          size instead of `None`. Note that in this case the graph (and the Keras
          layer) can only be used on a specific batch size. If you want to use a
          different batch size, you need to create another `AsKeras` object
          with a different `batch_size`.
        initializer_rng: (optional) an RNG key used to create the weights and
          state if `trax_layer` doesn't have them. If `None`,
          `trax.fastmath.random.get_prng(0)` will be used.
        rng: (optional) an RNG key for the forward function (aka the "forward
          key"). If `None`, `trax.fastmath.random.get_prng(0)` will be used.
        rng_updater: (optional) a function of type rng_key -> rng_key, used to
          update the forward key after each forward pass. If `None`, the function
          `lambda x: trax.fastmath.random.split(x, 1)[0]` will be used, which
          advances the RNG key.
        dtype: (optional) the dtype of the inputs. See the `dtype` argument of
          `tf.keras.layers.Layer.__init__` for details.
      """
      super().__init__(dtype=dtype)
      with math_lib.use_backend(math_lib.Backend.TFNP):
          if initializer_rng is None:
              initializer_rng = math_lib.random.get_prng(0)
          if rng is None:
              rng = math_lib.random.get_prng(0)
          if rng_updater is None:
              rng_updater = lambda x: math_lib.random.split(x, 1)[0]
          self._trax_layer = trax_layer
          self._batch_size = batch_size
          self._initializer_rng = initializer_rng
          self._forward_rng_init = rng
          self._rng_updater = rng_updater
  
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 1 to create a new AsKeras object from inherited base classes: tf.keras.layers with example: obj = AsKeras(trax_layer, batch_size, initializer_rng, rng, rng_updater, dtype)