def pure(self, val: Union[(TfVal, core.Unit)]) -> TensorFlowTracer:
    'Lifts a non-Tracer into the TensorFlowTracer.\n\n    This function may be called by way of trace.full_raise.\n\n    The value may be a core.unit. During JAX transformations we sometimes\n    produce a Jaxpr that has arguments of abstract value core.abstract_unit\n    and results equal to core.unit. These are arguments and results that are\n    not used in the computation.\n\n    In TF world, we represent core.unit as NaN. This is safe, as these values\n    should never be used.\n    '
    if (val is core.unit):
        return TensorFlowTracer(self, tf.constant(np.nan, tf.float32), core.abstract_unit)
    else:
        (shape, dtype) = _tfval_shape_dtype(val)
        return TensorFlowTracer(self, val, core.ShapedArray(shape, dtype))
Lifts a non-Tracer into the TensorFlowTracer.

This function may be called by way of trace.full_raise.

The value may be a core.unit. During JAX transformations we sometimes
produce a Jaxpr that has arguments of abstract value core.abstract_unit
and results equal to core.unit. These are arguments and results that are
not used in the computation.

In TF world, we represent core.unit as NaN. This is safe, as these values
should never be used.
def pure(self, val: Union[(TfVal, core.Unit)]) -> TensorFlowTracer:
    'Lifts a non-Tracer into the TensorFlowTracer.\n\n    This function may be called by way of trace.full_raise.\n\n    The value may be a core.unit. During JAX transformations we sometimes\n    produce a Jaxpr that has arguments of abstract value core.abstract_unit\n    and results equal to core.unit. These are arguments and results that are\n    not used in the computation.\n\n    In TF world, we represent core.unit as NaN. This is safe, as these values\n    should never be used.\n    '
    if (val is core.unit):
        return TensorFlowTracer(self, tf.constant(np.nan, tf.float32), core.abstract_unit)
    else:
        (shape, dtype) = _tfval_shape_dtype(val)
        return TensorFlowTracer(self, val, core.ShapedArray(shape, dtype))