we adapted the type promotion lattice from jax ref https github comgooglejaxblobmainjaxsrcdtypes py return the type promotion lattice in the form of a dag this dag maps each type to its immediately higher type on the lattice compute the least upper bound of a set of nodes args nodes sequence of entries from dtypes weaktypes returns the type representing the least upper bound of the input nodes on the promotion lattice this function computes the least upper bound of a set of nodes n within a partially ordered set defined by the lattice generated above given a partially ordered set s let the set of upper bounds of n s be ubn m s n m further for a set of nodes n s let the set of common upper bounds be given by cubn a s b n a ubb then the least upper bound of n is defined as lubn c cubn d cubn c d the definition of an upper bound implies that c d if and only if d ubc so the lub can be expressed lubn c cubn d cubn d ubc or equivalently lubn c cubn cubn ubc by definition lubn has a cardinality of 1 for a partially ordered set note a potential algorithmic shortcut from the definition of cubn we have c n cubn ubc so if n cubn is nonempty if follows that lubn n cubn if we get here it means the lattice is illformed return a dtype weaktype tuple for the given input isweaktype false if value is int or value is float note that we can t use value in int float because the dtype might be equal to python scalar types e g tf float32 float is true isweaktype true return standardizedtypevalue isweaktype functools lrucachemaxsizenone def respectweaktypedtype weaktype resolve weak type by the precision of backend floatx extendedalloweddtypes alloweddtypes unionweaktypes if dtype not in extendedalloweddtypes raise valueerror invalid value for argument dtype expected one of fextendedalloweddtypes received dtypedtype if precision not in 16 32 64 raise valueerror finvalid value for argument precision expected one of f 16 32 64 received precisionprecision if dtype bfloat16 special case for bfloat16 dtypeindicator f else dtypeindicator dtype 1 if dtypeindicator b return bool elif dtypeindicator i return int precision elif dtypeindicator u return uint precision else return float precision bit64tobit16dtype int32 int16 int64 int16 uint32 uint16 uint64 uint16 float32 float16 float64 float16 bit64tobit32dtype int64 int32 uint64 uint32 float64 float32 def latticeresulttypeargs dtypes weaktypes zipdtypeandweaktypearg for arg in args if lendtypes 1 outdtype dtypes0 outweaktype weaktypes0 elif lensetdtypes 1 and not allweaktypes trivial promotion case this allows extended dtypes through outdtype dtypes0 outweaktype false elif allweaktypes if all inputs are weakly typed we compute the bound of the stronglytyped counterparts and apply the weak type at the end this avoids returning the incorrect result with noncanonical weak types e g weak int16 outdtype leastupperbound respectweaktyped false for d in dtypes outweaktype true else outdtype leastupperbound respectweaktyped w for d w in zipdtypes weaktypes outweaktype anyoutdtype is t for t in weaktypes outweaktype outdtype bool and outweaktype precision backend floatx2 if outweaktype outdtype resolveweaktypeoutdtype precisionprecision return outdtype kerasexportkeras backend resulttype def resulttypedtypes if lendtypes 0 if no dtypes provided default to floatx this matches ops converttotensor return backend floatx return latticeresulttype backend floatx if arg is none else arg for arg in dtypes we adapted the type promotion lattice from jax ref https github com google jax blob main jax _src dtypes py return the type promotion lattice in the form of a dag this dag maps each type to its immediately higher type on the lattice compute the least upper bound of a set of nodes args nodes sequence of entries from dtypes weak_types returns the type representing the least upper bound of the input nodes on the promotion lattice this function computes the least upper bound of a set of nodes n within a partially ordered set defined by the lattice generated above given a partially ordered set s let the set of upper bounds of n s be ub n m s n m further for a set of nodes n s let the set of common upper bounds be given by cub n a s b n a ub b then the least upper bound of n is defined as lub n c cub n d cub n c d the definition of an upper bound implies that c d if and only if d ub c so the lub can be expressed lub n c cub n d cub n d ub c or equivalently lub n c cub n cub n ub c by definition lub n has a cardinality of 1 for a partially ordered set note a potential algorithmic shortcut from the definition of cub n we have c n cub n ub c so if n cub n is nonempty if follows that lub n n cub n if we get here it means the lattice is ill formed return a dtype weak_type tuple for the given input note that we can t use value in int float because the dtype might be equal to python scalar types e g tf float32 float is true return the weak dtype of dtype if weak_type true resolve weak type by the precision of backend floatx special case for bfloat16 trivial promotion case this allows extended dtypes through if all inputs are weakly typed we compute the bound of the strongly typed counterparts and apply the weak type at the end this avoids returning the incorrect result with non canonical weak types e g weak int16 returns the type from applying the keras type promotion rules in general each argument is first parsed by backend standardize_dtype and the resulting dtype is determined by the least upper bound of the type promotion lattice note this function attempts to match the result of jnp result_type args dtypes input dtypes returns the result dtype examples x keras ops ones 1 dtype bfloat16 keras backend result_type x dtype int bfloat16 x keras ops ones 1 dtype int32 y keras ops ones 1 dtype float32 keras backend result_type x dtype y dtype float32 if no dtypes provided default to floatx this matches ops convert_to_tensor
import functools

from keras import backend
from keras.api_export import keras_export
from keras.backend.common.variables import ALLOWED_DTYPES
from keras.backend.common.variables import standardize_dtype



BOOL_TYPES = ["bool"]
INT_TYPES = [
    "uint8",
    "uint16",
    "uint32",
    "uint64",
    "int8",
    "int16",
    "int32",
    "int64",
]
FLOAT_TYPES = ["bfloat16", "float16", "float32", "float64"]
WEAK_TYPES = ["int", "float"]


def _type_promotion_lattice():
    
    (b1,) = BOOL_TYPES
    (u1, u2, u4, u8, i1, i2, i4, i8) = INT_TYPES
    bf, f2, f4, f8 = FLOAT_TYPES
    i_, f_ = WEAK_TYPES
    out = {
        b1: [i_],
        u1: [i2, u2],
        u2: [i4, u4],
        u4: [i8, u8],
        u8: [f_],
        i_: [u1, i1],
        i1: [i2],
        i2: [i4],
        i4: [i8],
        i8: [f_],
        f_: [bf, f2],
        bf: [f4],
        f2: [f4],
        f4: [f8],
        f8: [],
    }
    return out


def _make_lattice_upper_bounds():
    lattice = _type_promotion_lattice()
    upper_bounds = {node: {node} for node in lattice}
    for n in lattice:
        while True:
            new_upper_bounds = set().union(
                *(lattice[b] for b in upper_bounds[n])
            )
            if n in new_upper_bounds:
                raise ValueError(
                    f"cycle detected in type promotion lattice for node {n}"
                )
            if new_upper_bounds.issubset(upper_bounds[n]):
                break
            upper_bounds[n] |= new_upper_bounds
    return upper_bounds


LATTICE_UPPER_BOUNDS = _make_lattice_upper_bounds()


@functools.lru_cache(512)
def _least_upper_bound(*nodes):
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    N = set(nodes)
    UB = LATTICE_UPPER_BOUNDS
    try:
        bounds = [UB[n] for n in N]
    except KeyError:
        dtype = next(n for n in N if n not in UB)
        raise ValueError(
            f"{dtype=} is not a valid dtype for Keras type promotion."
        )
    CUB = set.intersection(*bounds)
    LUB = (CUB & N) or {c for c in CUB if CUB.issubset(UB[c])}
    if len(LUB) == 1:
        return LUB.pop()
    elif len(LUB) == 0:
        msg = (
            f"Input dtypes {tuple(str(n) for n in nodes)} have no available "
            "implicit dtype promotion path. Try explicitly casting inputs to "
            "the desired output type."
        )
        raise ValueError(msg)
    else:
        
        raise ValueError(
            f"Internal Type Promotion error: {nodes} do not have a unique "
            f"least upper bound on the specified lattice; options are {LUB}. "
            "This is an unexpected error in Keras's internal logic; "
            "please report it to the maintainers."
        )


def _dtype_and_weaktype(value):
    
    is_weak_type = False
    if value is int or value is float:
        
        
        
        is_weak_type = True
    return standardize_dtype(value), is_weak_type


@functools.lru_cache(maxsize=None)
def _respect_weak_type(dtype, weak_type):
    
    if weak_type:
        if dtype == "bool":
            return dtype
        elif "float" in dtype:
            return "float"
        elif "int" in dtype:
            return "int"
        else:
            raise ValueError(
                "Invalid value for argument `dtype`. Expected one of "
                f"{ALLOWED_DTYPES}. Received: dtype={dtype}"
            )
    return dtype


@functools.lru_cache(maxsize=None)
def _resolve_weak_type(dtype, precision="32"):
    
    extended_allowed_dtypes = ALLOWED_DTYPES.union(WEAK_TYPES)
    if dtype not in extended_allowed_dtypes:
        raise ValueError(
            "Invalid value for argument `dtype`. Expected one of "
            f"{extended_allowed_dtypes}. Received: dtype={dtype}"
        )
    if precision not in ["16", "32", "64"]:
        raise ValueError(
            f"Invalid value for argument `precision`. Expected one of "
            f"('16', '32', '64'). Received: precision={precision}"
        )
    if dtype == "bfloat16":  
        dtype_indicator = "f"
    else:
        dtype_indicator = dtype[:1]

    if dtype_indicator == "b":
        return "bool"
    elif dtype_indicator == "i":
        return "int" + precision
    elif dtype_indicator == "u":
        return "uint" + precision
    else:
        return "float" + precision


BIT64_TO_BIT16_DTYPE = {
    "int32": "int16",
    "int64": "int16",
    "uint32": "uint16",
    "uint64": "uint16",
    "float32": "float16",
    "float64": "float16",
}
BIT64_TO_BIT32_DTYPE = {
    "int64": "int32",
    "uint64": "uint32",
    "float64": "float32",
}


def _lattice_result_type(*args):
    dtypes, weak_types = zip(*(_dtype_and_weaktype(arg) for arg in args))
    if len(dtypes) == 1:
        out_dtype = dtypes[0]
        out_weak_type = weak_types[0]
    elif len(set(dtypes)) == 1 and not all(weak_types):
        
        out_dtype = dtypes[0]
        out_weak_type = False
    elif all(weak_types):
        
        
        
        
        out_dtype = _least_upper_bound(
            *{_respect_weak_type(d, False) for d in dtypes}
        )
        out_weak_type = True
    else:
        out_dtype = _least_upper_bound(
            *{_respect_weak_type(d, w) for d, w in zip(dtypes, weak_types)}
        )
        out_weak_type = any(out_dtype is t for t in WEAK_TYPES)

    out_weak_type = (out_dtype != "bool") and out_weak_type
    precision = backend.floatx()[-2:]
    if out_weak_type:
        out_dtype = _resolve_weak_type(out_dtype, precision=precision)
    return out_dtype


@keras_export("keras.backend.result_type")
def result_type(*dtypes):
    
    if len(dtypes) == 0:
        
        
        return backend.floatx()
    return _lattice_result_type(
        *(backend.floatx() if arg is None else arg for arg in dtypes),
    )