def _dot_general(lhs, rhs, *, dimension_numbers, precision: Optional[Tuple[(PrecisionType, PrecisionType)]], preferred_element_type: Optional[DType], _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue):
    'Implementation of lax.dot_general_p in terms of tf.linalg.einsum.'
    ((lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch)) = dimension_numbers
    (lhs_ndim, rhs_ndim) = (len(lhs.shape), len(rhs.shape))
    if _enable_xla:
        dnums_proto = xla_data_pb2.DotDimensionNumbers()
        dnums_proto.lhs_contracting_dimensions.extend(lhs_contracting)
        dnums_proto.rhs_contracting_dimensions.extend(rhs_contracting)
        dnums_proto.lhs_batch_dimensions.extend(lhs_batch)
        dnums_proto.rhs_batch_dimensions.extend(rhs_batch)
        precision_config_proto = _precision_config_proto(precision)
        res = tfxla.dot_general(lhs, rhs, dnums_proto, precision_config_proto, preferred_element_type=preferred_element_type)
        res.set_shape(_aval_to_tf_shape(_out_aval))
        return res
    if ((lhs_batch == rhs_batch == tuple(range(len(lhs_batch)))) and ((lhs_ndim - rhs_ndim) in [(- 1), 0, 1]) and (1 <= (lhs_ndim - len(lhs_batch)) <= 2) and (1 <= (rhs_ndim - len(rhs_batch)) <= 2) and (lhs_contracting == ((len(lhs.shape) - 1),)) and (rhs_contracting == (len(lhs_batch),))):
        squeeze_idxs = []
        if ((lhs_ndim - len(lhs_batch)) == 1):
            lhs = tf.expand_dims(lhs, (lhs_ndim - 1))
            squeeze_idxs.append((len(lhs.shape) - 2))
        if ((rhs_ndim - len(rhs_batch)) == 1):
            rhs = tf.expand_dims(rhs, rhs_ndim)
            squeeze_idxs.append((len(rhs.shape) - 1))
        result = tf.linalg.matmul(lhs, rhs)
        if (len(squeeze_idxs) != 0):
            assert all([(result.shape[i] == 1) for i in squeeze_idxs])
            result = tf.squeeze(result, squeeze_idxs)
        return result
    new_id = iter(string.ascii_letters)
    lhs_axis_ids = [next(new_id) for _ in lhs.shape]
    rhs_axis_ids = [next(new_id) for _ in rhs.shape]
    lhs_out_axis_ids = lhs_axis_ids[:]
    rhs_out_axis_ids = rhs_axis_ids[:]
    for (lhs_axis, rhs_axis) in zip(lhs_contracting, rhs_contracting):
        shared_id = next(new_id)
        lhs_axis_ids[lhs_axis] = shared_id
        rhs_axis_ids[rhs_axis] = shared_id
        lhs_out_axis_ids[lhs_axis] = None
        rhs_out_axis_ids[rhs_axis] = None
    batch_ids = []
    for (lhs_axis, rhs_axis) in zip(lhs_batch, rhs_batch):
        shared_id = next(new_id)
        lhs_axis_ids[lhs_axis] = shared_id
        rhs_axis_ids[rhs_axis] = shared_id
        lhs_out_axis_ids[lhs_axis] = None
        rhs_out_axis_ids[rhs_axis] = None
        batch_ids.append(shared_id)
    not_none = (lambda x: (x is not None))
    out_axis_ids = list(filter(not_none, ((batch_ids + lhs_out_axis_ids) + rhs_out_axis_ids)))
    assert (lhs.dtype == rhs.dtype)
    spec = '{},{}->{}'.format(''.join(lhs_axis_ids), ''.join(rhs_axis_ids), ''.join(out_axis_ids))
    return tf.linalg.einsum(spec, lhs, rhs)
Implementation of lax.dot_general_p in terms of tf.linalg.einsum.
def _dot_general(lhs, rhs, *, dimension_numbers, precision: Optional[Tuple[(PrecisionType, PrecisionType)]], preferred_element_type: Optional[DType], _in_avals: Sequence[core.AbstractValue], _out_aval: core.AbstractValue):
    
    ((lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch)) = dimension_numbers
    (lhs_ndim, rhs_ndim) = (len(lhs.shape), len(rhs.shape))
    if _enable_xla:
        dnums_proto = xla_data_pb2.DotDimensionNumbers()
        dnums_proto.lhs_contracting_dimensions.extend(lhs_contracting)
        dnums_proto.rhs_contracting_dimensions.extend(rhs_contracting)
        dnums_proto.lhs_batch_dimensions.extend(lhs_batch)
        dnums_proto.rhs_batch_dimensions.extend(rhs_batch)
        precision_config_proto = _precision_config_proto(precision)
        res = tfxla.dot_general(lhs, rhs, dnums_proto, precision_config_proto, preferred_element_type=preferred_element_type)
        res.set_shape(_aval_to_tf_shape(_out_aval))
        return res
    if ((lhs_batch == rhs_batch == tuple(range(len(lhs_batch)))) and ((lhs_ndim - rhs_ndim) in [(- 1), 0, 1]) and (1 <= (lhs_ndim - len(lhs_batch)) <= 2) and (1 <= (rhs_ndim - len(rhs_batch)) <= 2) and (lhs_contracting == ((len(lhs.shape) - 1),)) and (rhs_contracting == (len(lhs_batch),))):
        squeeze_idxs = []
        if ((lhs_ndim - len(lhs_batch)) == 1):
            lhs = tf.expand_dims(lhs, (lhs_ndim - 1))
            squeeze_idxs.append((len(lhs.shape) - 2))
        if ((rhs_ndim - len(rhs_batch)) == 1):
            rhs = tf.expand_dims(rhs, rhs_ndim)
            squeeze_idxs.append((len(rhs.shape) - 1))
        result = tf.linalg.matmul(lhs, rhs)
        if (len(squeeze_idxs) != 0):
            assert all([(result.shape[i] == 1) for i in squeeze_idxs])
            result = tf.squeeze(result, squeeze_idxs)
        return result
    new_id = iter(string.ascii_letters)
    lhs_axis_ids = [next(new_id) for _ in lhs.shape]
    rhs_axis_ids = [next(new_id) for _ in rhs.shape]
    lhs_out_axis_ids = lhs_axis_ids[:]
    rhs_out_axis_ids = rhs_axis_ids[:]
    for (lhs_axis, rhs_axis) in zip(lhs_contracting, rhs_contracting):
        shared_id = next(new_id)
        lhs_axis_ids[lhs_axis] = shared_id
        rhs_axis_ids[rhs_axis] = shared_id
        lhs_out_axis_ids[lhs_axis] = None
        rhs_out_axis_ids[rhs_axis] = None
    batch_ids = []
    for (lhs_axis, rhs_axis) in zip(lhs_batch, rhs_batch):
        shared_id = next(new_id)
        lhs_axis_ids[lhs_axis] = shared_id
        rhs_axis_ids[rhs_axis] = shared_id
        lhs_out_axis_ids[lhs_axis] = None
        rhs_out_axis_ids[rhs_axis] = None
        batch_ids.append(shared_id)
    not_none = (lambda x: (x is not None))
    out_axis_ids = list(filter(not_none, ((batch_ids + lhs_out_axis_ids) + rhs_out_axis_ids)))
    assert (lhs.dtype == rhs.dtype)
    spec = '{},{}->{}'.format(.join(lhs_axis_ids), .join(rhs_axis_ids), .join(out_axis_ids))
    return tf.linalg.einsum(spec, lhs, rhs)