code: def _ConvertToNaNsOnAnyZero():
      def _convert_to_nans(x, y):
          # if all values in y are non-zeros, return x; otherwise return 0s
          return jnp.where(jnp.all(y, keepdims=False), x, x / 0.0), y
  
      return tl.Fn("ConvertToNaNsOnAnyZero", _convert_to_nans, n_out=2)
  
imports: import functools
  from trax import layers
  from trax.fastmath import numpy
  from trax.models.reformer import reformer
  from trax.models.research import configurable_transformer
  from trax.models.research import transformer2
label_desc: Use image node_id 7 for calling a global function with example usage: _ConvertToNaNsOnAnyZero() and returns: tl, jnp,  y