def __init__(self, trace: 'TensorFlowTrace', val: TfVal, aval: core.AbstractValue):
    self._trace = trace
    self._aval = aval
    if (aval is core.abstract_unit):
        self.val = val
    elif isinstance(val, (tf.Tensor, tf.Variable)):
        (val_shape, val_dtype) = _tfval_shape_dtype(val)
        aval_dtype = np.dtype(self._aval.dtype)
        if ((val_dtype != aval_dtype) and (not config.x64_enabled) and (((val_dtype == tf.int32) and (aval_dtype == jnp.int64)) or ((val_dtype == tf.int64) and (aval_dtype == jnp.int32)) or ((val_dtype == tf.float32) and (aval_dtype == jnp.float64)) or ((val_dtype == tf.float64) and (aval_dtype == jnp.float32)) or ((val_dtype == tf.complex128) and (aval_dtype == jnp.complex64)))):
            val = tf.cast(val, dtype=aval_dtype)
            val_dtype = aval_dtype
        if config.jax_enable_checks:
            assert (aval_dtype == val_dtype), f'expected {aval_dtype} == {val_dtype}'
            for (aval_dim, val_dim) in util.safe_zip(self._aval.shape, val_shape):
                if (val_dim is None):
                    assert isinstance(aval_dim, shape_poly.DimVar), f'expected {self._aval.shape} == {val_shape}'
                elif (not isinstance(aval_dim, shape_poly.DimVar)):
                    assert (aval_dim == val_dim), f'expected {self._aval.shape} == {val_shape}'
                else:
                    try:
                        aval_int = int(_eval_shape([aval_dim]))
                    except TypeError:
                        continue
                    assert (aval_int == val_dim), f'expected {self._aval.shape} == {val_shape}. Found {aval_int} != {val_dim}.'
        self.val = val
    else:
        self.val = _safe_convert_to_tensor(val, dtype=self._aval.dtype)
def __init__(self, trace: 'TensorFlowTrace', val: TfVal, aval: core.AbstractValue):
    self._trace = trace
    self._aval = aval
    if (aval is core.abstract_unit):
        self.val = val
    elif isinstance(val, (tf.Tensor, tf.Variable)):
        (val_shape, val_dtype) = _tfval_shape_dtype(val)
        aval_dtype = np.dtype(self._aval.dtype)
        if ((val_dtype != aval_dtype) and (not config.x64_enabled) and (((val_dtype == tf.int32) and (aval_dtype == jnp.int64)) or ((val_dtype == tf.int64) and (aval_dtype == jnp.int32)) or ((val_dtype == tf.float32) and (aval_dtype == jnp.float64)) or ((val_dtype == tf.float64) and (aval_dtype == jnp.float32)) or ((val_dtype == tf.complex128) and (aval_dtype == jnp.complex64)))):
            val = tf.cast(val, dtype=aval_dtype)
            val_dtype = aval_dtype
        if config.jax_enable_checks:
            assert (aval_dtype == val_dtype), f'expected {aval_dtype} == {val_dtype}'
            for (aval_dim, val_dim) in util.safe_zip(self._aval.shape, val_shape):
                if (val_dim is None):
                    assert isinstance(aval_dim, shape_poly.DimVar), f'expected {self._aval.shape} == {val_shape}'
                elif (not isinstance(aval_dim, shape_poly.DimVar)):
                    assert (aval_dim == val_dim), f'expected {self._aval.shape} == {val_shape}'
                else:
                    try:
                        aval_int = int(_eval_shape([aval_dim]))
                    except TypeError:
                        continue
                    assert (aval_int == val_dim), f'expected {self._aval.shape} == {val_shape}. Found {aval_int} != {val_dim}.'
        self.val = val
    else:
        self.val = _safe_convert_to_tensor(val, dtype=self._aval.dtype)