class_docstr: Caches a weak chain of transformations for all bijector instances.
  
  On a cache miss, the other side of the transformation is produced
  and both sides are cached before returning a result. The input object
  is cached with a strong reference, while the transformed output is
  held weakly. When there are no more references to the result, both
  directions of the mapping are cleared simultaneously.
  
  The general contracts are:
    1) The inputs used on cache misses will always be retained
       until their corresponding outputs are abandoned.
    2) `forward` will return a cached result if and only if
       `inverse` would have returned a cached result.
  
  Retaining results in such a manner guarantees that there will be no
  circular references, allowing the garbage collector to remove "dead"
  links from the cache.
  
  Example usage:
  
  ```py
  bij = MyBijector()
  cache = BijectorCache(bijector=bij)
  
  x = tf.constant(1.)
  ref = weakref.ref(x)
  
  for i in range(10):
    assert (len(cache.weak_keys(direction='forward'))
            == len(cache.weak_keys(direction='inverse'))
            == i)
    x = cache.forward(x)
  
  for i in range(10):
    assert len(cache.weak_keys(direction='inverse')) == 10 - i
    x = cache.inverse(x)
    assert len(cache.weak_keys(direction='forward')) == 10 - i
  
  assert len(cache.weak_keys(direction='inverse')) == 0
  assert ref() == x
  
  x = None
  assert len(cache.weak_keys(direction='forward')) == 0
  ```
class_docstr_tok: ["Caches","a","weak","chain","of","transformations","for","all","bijector","instances",".","On","a","cache","miss",",","the","other","side","of","the","transformation","is","produced","and","both","sides","are","cached","before","returning","a","result",".","The","input","object","is","cached","with","a","strong","reference",",","while","the","transformed","output","is","held","weakly",".","When","there","are","no","more","references","to","the","result",",","both","directions","of","the","mapping","are","cleared","simultaneously",".","The","general","contracts","are",":","1",")","The","inputs","used","on","cache","misses","will","always","be","retained","until","their","corresponding","outputs","are","abandoned",".","2",")","`","forward","`","will","return","a","cached","result","if","and","only","if","`","inverse","`","would","have","returned","a","cached","result",".","Retaining","results","in","such","a","manner","guarantees","that","there","will","be","no","circular","references",",","allowing","the","garbage","collector","to","remove","``","dead","''","links","from","the","cache",".","Example","usage",":","``","`","py","bij","=","MyBijector","(",")","cache","=","BijectorCache","(","bijector=bij",")","x","=","tf.constant","(","1",".",")","ref","=","weakref.ref","(","x",")","for","i","in","range","(","10",")",":","assert","(","len","(","cache.weak_keys","(","direction='forward","'",")",")","==","len","(","cache.weak_keys","(","direction='inverse","'",")",")","==","i",")","x","=","cache.forward","(","x",")","for","i","in","range","(","10",")",":","assert","len","(","cache.weak_keys","(","direction='inverse","'",")",")","==","10","-","i","x","=","cache.inverse","(","x",")","assert","len","(","cache.weak_keys","(","direction='forward","'",")",")","==","10","-","i","assert","len","(","cache.weak_keys","(","direction='inverse","'",")",")","==","0","assert","ref","(",")","==","x","x","=","None","assert","len","(","cache.weak_keys","(","direction='forward","'",")",")","==","0","``","`"]
docstr: Returns the keys in input_cache.
docstr_tok: ["Returns","the","keys","in","input_cache","."]
code: def weak_keys(
      self, bijector=None, bijector_class=None, direction=None
  ):
      """Returns the keys in input_cache."""
      bijector = bijector or self.bijector
      bijector_class = bijector_class or self.bijector_class
  
      if direction == "forward":
          fn_name = self._forward_name
      elif direction == "inverse":
          fn_name = self._inverse_name
      elif direction is None:
          fn_name = None
      else:
          raise ValueError(
              '`direction` must be `"forward"`, `"inverse"`, or '
              "`None`. Saw: {}".format(direction)
          )
      out = []
      # Make a copy of the keys, in case one of them gets deleted, so as to avoid
      # the dict-size changing at runtime error.
      for k in list(self.storage.keys()):
          bijector_key, bijector_class_key, fn_key, _ = k.subkey
          if (
              k.alive
              and (bijector is None or bijector_key == bijector)
              and (
                  bijector_class is None
                  or bijector_class == bijector_class_key
              )
              and (fn_name is None or fn_key == fn_name)
          ):
              out.append(k)
      return out
  
imports: import collections
  import weakref
  import numpy
  import tensorflow.compat.v2
  from tensorflow.python.util import nest
label_desc: Use image node_id 9 for calling the BijectorCache obj's underlying member method code with example usage: obj.weak_keys(bijector, bijector_class, direction) and returns: out