def process_primitive(self, primitive: core.Primitive, tracers: Sequence[TensorFlowTracer], params) -> TensorFlowTracer:
    (impl, impl_needs_avals) = self.get_primitive_impl(primitive)
    args_avals: Sequence[core.AbstractValue] = tuple((t.aval for t in tracers))
    out_aval = primitive.abstract_eval(*args_avals, **params)
    args_tf: Sequence[TfVal] = [t.val for t in tracers]
    if impl_needs_avals:
        val_out: TfVal = impl(*args_tf, _in_avals=args_avals, _out_aval=out_aval, **params)
    else:
        val_out = impl(*args_tf, **params)
    if primitive.multiple_results:
        out = [TensorFlowTracer(self, v, a) for (v, a) in util.safe_zip(val_out, out_aval)]
    else:
        out = TensorFlowTracer(self, val_out, out_aval)
    if config.jax_enable_checks:
        if primitive.multiple_results:
            for (o, expected_aval) in zip(out, out_aval):
                assert (o.aval.strip_weak_type() == expected_aval.strip_weak_type()), f'{primitive}: out.aval = {o.aval}; expected {expected_aval}'
        else:
            assert (out.aval == out_aval), f'{primitive}: out.aval = {out.aval}; expected {out_aval}'
    return out
def process_primitive(self, primitive: core.Primitive, tracers: Sequence[TensorFlowTracer], params) -> TensorFlowTracer:
    (impl, impl_needs_avals) = self.get_primitive_impl(primitive)
    args_avals: Sequence[core.AbstractValue] = tuple((t.aval for t in tracers))
    out_aval = primitive.abstract_eval(*args_avals, **params)
    args_tf: Sequence[TfVal] = [t.val for t in tracers]
    if impl_needs_avals:
        val_out: TfVal = impl(*args_tf, _in_avals=args_avals, _out_aval=out_aval, **params)
    else:
        val_out = impl(*args_tf, **params)
    if primitive.multiple_results:
        out = [TensorFlowTracer(self, v, a) for (v, a) in util.safe_zip(val_out, out_aval)]
    else:
        out = TensorFlowTracer(self, val_out, out_aval)
    if config.jax_enable_checks:
        if primitive.multiple_results:
            for (o, expected_aval) in zip(out, out_aval):
                assert (o.aval.strip_weak_type() == expected_aval.strip_weak_type()), f'{primitive}: out.aval = {o.aval}; expected {expected_aval}'
        else:
            assert (out.aval == out_aval), f'{primitive}: out.aval = {out.aval}; expected {out_aval}'
    return out